fix(tf): fix DeepEval degradation for virtual types

[njzjz]

The energy fitting network computes nloc by summing over natoms[2:] instead of reading natoms[0].

deepmd-kit/deepmd/tf/fit/ener.py

Lines 717 to 720 in 8dab33b

	self.atom_bias_ener = tf.reshape(
	tf.nn.embedding_lookup(self.t_bias_atom_e, self.atype_nloc),
	[tf.shape(inputs)[0], tf.reduce_sum(natoms[2 : 2 + ntypes_atom])],
	)

This causes an issue for the virtual types after refactoring DeepEval. Before, natoms_vec is [nloc, nall, nloc, ...] for mixed types. After refactoring, we use the same natoms_vec for mixed_types and the normal case, so the virtual type support is broken.

This was not detected by the test, as the test model for the virtual types was added 12 months ago, but the energy fitting was changed 10 months ago.

[codecov]

Codecov Report

All modified and coverable lines are covered by tests ✅

Project coverage is 77.35%. Comparing base (487f85c) to head (8dab33b).

Additional details and impacted files

@@           Coverage Diff           @@
##            devel    #3464   +/-   ##
=======================================
  Coverage   77.35%   77.35%           
=======================================
  Files         428      428           
  Lines       37000    37004    +4     
  Branches     1620     1620           
=======================================
+ Hits        28622    28626    +4     
  Misses       7510     7510           
  Partials      868      868           

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