[Code scan] Move electronic-configuration type embeddings onto the active backend

[njzjz]

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

Problem

Electronic-configuration type embeddings remain NumPy arrays even when the type-embedding network parameters are on another array backend.

Evidence:

• TypeEmbedNet.__init__() stores self.econf_tebd from get_econf_tebd() when use_econf_tebd=True:

  deepmd-kit/deepmd/dpmodel/utils/type_embed.py

  Lines 70 to 86 in 73de44b

  	) -> None:
  	self.ntypes = ntypes
  	self.neuron = neuron
  	self.seed = seed
  	self.resnet_dt = resnet_dt
  	self.precision = precision
  	self.activation_function = str(activation_function)
  	self.trainable = trainable
  	self.padding = padding
  	self.use_econf_tebd = use_econf_tebd
  	self.use_tebd_bias = use_tebd_bias
  	self.type_map = type_map
  	embed_input_dim = ntypes
  	if self.use_econf_tebd:
  	self.econf_tebd, embed_input_dim = get_econf_tebd(
  	self.type_map, precision=self.precision
  	)

• TypeEmbedNet.call() selects xp from the sample network weight, but the electronic-configuration branch passes self.econf_tebd directly into the embedding network:

  deepmd-kit/deepmd/dpmodel/utils/type_embed.py

  Lines 98 to 112 in 73de44b

  	def call(self) -> Array:
  	"""Compute the type embedding network."""
  	sample_array = self.embedding_net[0]["w"]
  	xp = array_api_compat.array_namespace(sample_array)
  	if not self.use_econf_tebd:
  	embed = self.embedding_net(
  	xp.eye(
  	self.ntypes,
  	dtype=sample_array.dtype,
  	device=array_api_compat.device(sample_array),
  	)
  	)
  	else:
  	embed = self.embedding_net(self.econf_tebd)
  	if self.padding:

• get_econf_tebd() always constructs a NumPy array:

  deepmd-kit/deepmd/dpmodel/utils/type_embed.py

  Lines 248 to 273 in 73de44b

  	def get_econf_tebd(
  	type_map: list[str], precision: str = "default"
  	) -> tuple[Array, int]:
  	from deepmd.utils.econf_embd import (
  	ECONF_DIM,
  	)
  	from deepmd.utils.econf_embd import (
  	normalized_electronic_configuration_embedding as electronic_configuration_embedding,
  	)
  	from deepmd.utils.econf_embd import type_map as periodic_table
  	assert type_map is not None, (
  	"When using electronic configuration type embedding, type_map must be provided!"
  	)
  	missing_types = [t for t in type_map if t not in periodic_table]
  	assert not missing_types, (
  	"When using electronic configuration type embedding, "
  	"all element in type_map should be in periodic table! "
  	f"Found these invalid elements: {missing_types}"
  	)
  	econf_tebd = np.array(
  	[electronic_configuration_embedding[kk] for kk in type_map],
  	dtype=PRECISION_DICT[precision],
  	)
  	embed_input_dim = ECONF_DIM

• NativeLayer.call() chooses its array namespace from the input x, then converts weights with that namespace:

  deepmd-kit/deepmd/dpmodel/utils/network.py

  Lines 268 to 303 in 73de44b

  	def call(self, x): # noqa: ANN001, ANN201
  	"""Forward pass.
  	Parameters
  	----------
  	x : Array
  	The input.
  	Returns
  	-------
  	np.ndarray
  	The output.
  	"""
  	if self.w is None or self.activation_function is None:
  	raise ValueError("w, b, and activation_function must be set")
  	xp = array_api_compat.array_namespace(x)
  	fn = get_activation_fn(self.activation_function)
  	device = array_api_compat.device(x)
  	w = xp.asarray(self.w[...], device=device)
  	y = (
  	xp.matmul(x, w) + xp.asarray(self.b[...], device=device)
  	if self.b is not None
  	else xp.matmul(x, w)
  	)
  	if y.dtype != x.dtype:
  	# workaround for bfloat16
  	# https://github.com/jax-ml/ml_dtypes/issues/235
  	y = xp.astype(y, x.dtype)
  	y = fn(y)
  	if self.idt is not None:
  	y = y * xp.asarray(self.idt, device=array_api_compat.device(x))
  	if self.resnet and self.w.shape[1] == self.w.shape[0]:
  	y = y + x
  	elif self.resnet and self.w.shape[1] == 2 * self.w.shape[0]:
  	y = y + xp.concat([x, x], axis=-1)
  	return y

Impact

When weights are torch/JAX arrays but use_econf_tebd=True, the embedding network can run under the NumPy namespace, forcing backend weights through NumPy conversion or failing for non-CPU/device-backed arrays. Downstream descriptor code that expects backend-native type embeddings can then fail on namespace mismatches.

Suggested Fix

Convert self.econf_tebd to the active namespace, dtype, and device before passing it to the embedding network. Add a torch or array-api-strict test for TypeEmbedNet(use_econf_tebd=True) with backend-native weights.