Fix a potential slice bug in se_t descriptor

[denghuilu]

In the code review process, I found that env_j in se_t type always start the slice from position 0 while the slice size evaluates from self.sel_a[type_j] due to the setting of the start_index_j:

deepmd-kit/deepmd/descriptor/se_t.py

Lines 462 to 498 in 1a25414

	with tf.variable_scope(name, reuse=reuse):
	start_index_i = 0
	result = None
	for type_i in range(self.ntypes):
	# cut-out inputs
	# with natom x (nei_type_i x 4)
	inputs_i = tf.slice (inputs,
	[ 0, start_index_i *4],
	[-1, self.sel_a[type_i] *4] )
	start_index_i += self.sel_a[type_i]
	nei_type_i = self.sel_a[type_i]
	shape_i = inputs_i.get_shape().as_list()
	assert(shape_i[1] == nei_type_i * 4)
	# with natom x nei_type_i x 4
	env_i = tf.reshape(inputs_i, [-1, nei_type_i, 4])
	# with natom x nei_type_i x 3
	env_i = tf.slice(env_i, [0, 0, 1], [-1, -1, -1])
	start_index_j = 0
	for type_j in range(type_i, self.ntypes):
	# with natom x (nei_type_j x 4)
	inputs_j = tf.slice (inputs,
	[ 0, start_index_j *4],
	[-1, self.sel_a[type_j] *4] )
	start_index_j += self.sel_a[type_j]
	nei_type_j = self.sel_a[type_j]
	shape_j = inputs_j.get_shape().as_list()
	assert(shape_j[1] == nei_type_j * 4)
	# with natom x nei_type_j x 4
	env_j = tf.reshape(inputs_j, [-1, nei_type_j, 4])
	# with natom x nei_type_i x 3
	env_j = tf.slice(env_j, [0, 0, 1], [-1, -1, -1])
	# with natom x nei_type_i x nei_type_j
	env_ij = tf.einsum('ijm,ikm->ijk', env_i, env_j)
	# with (natom x nei_type_i x nei_type_j)
	ebd_env_ij = tf.reshape(env_ij, [-1, 1])
	# with (natom x nei_type_i x nei_type_j) x out_size
	ebd_env_ij = embedding_net(ebd_env_ij,

And after the talk with Wanrun, we found that this may lead to repeated values ​​of env_j when there are multiple types of atoms. So I change the start_index_j value to avoid the replication of env_j. @amcadmus Please confirm the necessity of this change.

Note that this PR will not affect the single-atom type system.

[codecov-commenter]

Codecov Report

All modified and coverable lines are covered by tests ✅

Project coverage is 75.74%. Comparing base (c824ff6) to head (544cc69).
Report is 1469 commits behind head on devel.

Additional details and impacted files

@@           Coverage Diff           @@
##            devel    #1087   +/-   ##
=======================================
  Coverage   75.74%   75.74%           
=======================================
  Files          88       88           
  Lines        6996     6996           
=======================================
  Hits         5299     5299           
  Misses       1697     1697           

☔ View full report in Codecov by Sentry.
📢 Have feedback on the report? Share it here.