Found during a Codex global scan of deepmodeling/deepmd-kit at commit 73de44b1f94471b2e3bdb6b11f57b34d7bc791bb.

Problem

GeneralFitting initializes several parameter/statistic buffers as NumPy arrays and later passes them directly into operations using the input descriptor's array namespace.

Evidence:

• Buffers such as bias_atom_e, fparam_avg, aparam_avg, case_embd, and default_fparam_tensor are initialized as NumPy arrays:

  deepmd-kit/deepmd/dpmodel/fitting/general_fitting.py

  Lines 176 to 207 in 73de44b

  	net_dim_out = self._net_out_dim()
  	# init constants
  	if bias_atom_e is None:
  	self.bias_atom_e = np.zeros(
  	[self.ntypes, net_dim_out], dtype=GLOBAL_NP_FLOAT_PRECISION
  	)
  	else:
  	assert bias_atom_e.shape == (self.ntypes, net_dim_out)
  	self.bias_atom_e = bias_atom_e.astype(GLOBAL_NP_FLOAT_PRECISION)
  	if self.numb_fparam > 0:
  	self.fparam_avg = np.zeros(self.numb_fparam, dtype=self.prec)
  	self.fparam_inv_std = np.ones(self.numb_fparam, dtype=self.prec)
  	else:
  	self.fparam_avg, self.fparam_inv_std = None, None
  	if self.numb_aparam > 0:
  	self.aparam_avg = np.zeros(self.numb_aparam, dtype=self.prec)
  	self.aparam_inv_std = np.ones(self.numb_aparam, dtype=self.prec)
  	else:
  	self.aparam_avg, self.aparam_inv_std = None, None
  	if self.dim_case_embd > 0:
  	self.case_embd = np.zeros(self.dim_case_embd, dtype=self.prec)
  	else:
  	self.case_embd = None
  	if self.default_fparam is not None:
  	if self.numb_fparam > 0:
  	assert len(self.default_fparam) == self.numb_fparam, (
  	"default_fparam length mismatch!"
  	)
  	self.default_fparam_tensor = np.array(self.default_fparam, dtype=self.prec)
  	else:
  	self.default_fparam_tensor = None

• The forward path selects xp from the input descriptor and atom types:

  deepmd-kit/deepmd/dpmodel/fitting/general_fitting.py

  Lines 646 to 655 in 73de44b

  	xp = array_api_compat.array_namespace(descriptor, atype)
  	nf, nloc, nd = descriptor.shape
  	net_dim_out = self._net_out_dim()
  	# check input dim
  	if nd != self.dim_descrpt:
  	raise ValueError(
  	"get an input descriptor of dim {nd},"
  	"which is not consistent with {self.dim_descrpt}."
  	)
  	xx = descriptor

• It then uses xp.reshape/xp.tile on self.default_fparam_tensor and self.case_embd:

  deepmd-kit/deepmd/dpmodel/fitting/general_fitting.py

  Lines 666 to 721 in 73de44b

  	if self.numb_fparam > 0 and fparam is None:
  	# use default fparam
  	assert self.default_fparam_tensor is not None
  	fparam = xp.tile(
  	xp.reshape(self.default_fparam_tensor, (1, self.numb_fparam)), (nf, 1)
  	)
  	# check fparam dim, concate to input descriptor
  	if self.numb_fparam > 0:
  	assert fparam is not None, "fparam should not be None"
  	try:
  	fparam = xp.reshape(fparam, (nf, self.numb_fparam))
  	except (ValueError, RuntimeError) as e:
  	raise ValueError(
  	f"input fparam: cannot reshape {fparam.shape} "
  	f"into ({nf}, {self.numb_fparam})."
  	) from e
  	fparam = (fparam - self.fparam_avg[...]) * self.fparam_inv_std[...]
  	fparam = xp.tile(
  	xp.reshape(fparam, (nf, 1, self.numb_fparam)), (1, nloc, 1)
  	)
  	xx = xp.concat(
  	[xx, fparam],
  	axis=-1,
  	)
  	if xx_zeros is not None:
  	xx_zeros = xp.concat(
  	[xx_zeros, fparam],
  	axis=-1,
  	)
  	# check aparam dim, concate to input descriptor
  	if self.numb_aparam > 0 and not self.use_aparam_as_mask:
  	assert aparam is not None, "aparam should not be None"
  	try:
  	aparam = xp.reshape(aparam, (nf, nloc, self.numb_aparam))
  	except (ValueError, RuntimeError) as e:
  	raise ValueError(
  	f"input aparam: cannot reshape {aparam.shape} "
  	f"into ({nf}, {nloc}, {self.numb_aparam})."
  	) from e
  	aparam = (aparam - self.aparam_avg[...]) * self.aparam_inv_std[...]
  	xx = xp.concat(
  	[xx, aparam],
  	axis=-1,
  	)
  	if xx_zeros is not None:
  	xx_zeros = xp.concat(
  	[xx_zeros, aparam],
  	axis=-1,
  	)
  	if self.dim_case_embd > 0:
  	assert self.case_embd is not None
  	case_embd = xp.tile(
  	xp.reshape(self.case_embd[...], (1, 1, -1)), (nf, nloc, 1)
  	)

• It performs backend arithmetic with self.fparam_avg and self.aparam_avg:

  deepmd-kit/deepmd/dpmodel/fitting/general_fitting.py

  Lines 673 to 706 in 73de44b

  	# check fparam dim, concate to input descriptor
  	if self.numb_fparam > 0:
  	assert fparam is not None, "fparam should not be None"
  	try:
  	fparam = xp.reshape(fparam, (nf, self.numb_fparam))
  	except (ValueError, RuntimeError) as e:
  	raise ValueError(
  	f"input fparam: cannot reshape {fparam.shape} "
  	f"into ({nf}, {self.numb_fparam})."
  	) from e
  	fparam = (fparam - self.fparam_avg[...]) * self.fparam_inv_std[...]
  	fparam = xp.tile(
  	xp.reshape(fparam, (nf, 1, self.numb_fparam)), (1, nloc, 1)
  	)
  	xx = xp.concat(
  	[xx, fparam],
  	axis=-1,
  	)
  	if xx_zeros is not None:
  	xx_zeros = xp.concat(
  	[xx_zeros, fparam],
  	axis=-1,
  	)
  	# check aparam dim, concate to input descriptor
  	if self.numb_aparam > 0 and not self.use_aparam_as_mask:
  	assert aparam is not None, "aparam should not be None"
  	try:
  	aparam = xp.reshape(aparam, (nf, nloc, self.numb_aparam))
  	except (ValueError, RuntimeError) as e:
  	raise ValueError(
  	f"input aparam: cannot reshape {aparam.shape} "
  	f"into ({nf}, {nloc}, {self.numb_aparam})."
  	) from e
  	aparam = (aparam - self.aparam_avg[...]) * self.aparam_inv_std[...]

• It calls xp.astype(self.bias_atom_e[...], outs.dtype):

  deepmd-kit/deepmd/dpmodel/fitting/general_fitting.py

  Lines 776 to 783 in 73de44b

  	outs += xp.reshape(
  	xp.take(
  	xp.astype(self.bias_atom_e[...], outs.dtype),
  	xp.reshape(atype, (-1,)),
  	axis=0,
  	),
  	(nf, nloc, net_dim_out),
  	)

External torch checks show array_api_compat.torch.reshape(np_array, ...) and array_api_compat.torch.astype(np_array, ...) fail because the input is not a torch tensor.

Impact

Torch/JAX/array-API fitting paths can fail when defaults, case embeddings, or atomic energy biases are used, even if the descriptor and atom types are backend-native arrays.

Suggested Fix

Materialize these buffers on the active namespace/device before using them in xp.* operations, following the conversion pattern used by network layers. Add a fitting forward test with backend tensor descriptors, default fparam, nonzero dim_case_embd, and nonzero bias_atom_e.