From a0950550805bfabddd5768843eb92a13a4b13fb3 Mon Sep 17 00:00:00 2001 From: njzjz-bot Date: Fri, 17 Jul 2026 03:13:28 +0800 Subject: [PATCH] fix(dpmodel): move econf embeddings to active backend Materialize portable NumPy electronic-configuration tables in the namespace, dtype, and device of the type-embedding weights before evaluation. This prevents NativeLayer from selecting NumPy and coercing backend-native parameters back to CPU. Add a direct mixed-state regression covering Torch output identity, numerical parity, and change_type_map regeneration without relying on eager-converting backend wrappers. Coding-Agent: Codex Codex-Version: codex-cli 0.144.4 Model: gpt-5.6-sol Reasoning-Effort: xhigh --- deepmd/dpmodel/utils/type_embed.py | 11 +++- .../tests/consistent/test_type_embedding.py | 50 +++++++++++++++++++ 2 files changed, 60 insertions(+), 1 deletion(-) diff --git a/deepmd/dpmodel/utils/type_embed.py b/deepmd/dpmodel/utils/type_embed.py index f945d0f801..9d5ed78f46 100644 --- a/deepmd/dpmodel/utils/type_embed.py +++ b/deepmd/dpmodel/utils/type_embed.py @@ -115,7 +115,16 @@ def call(self) -> Array: ) ) else: - embed = self.embedding_net(self.econf_tebd) + # Electronic configurations are stored as NumPy for portable + # serialization. Materialize them beside the network parameters + # at call time so NativeLayer selects the active array backend + # instead of accidentally pulling device-backed weights to NumPy. + backend_econf_tebd = xp.asarray( + self.econf_tebd, + dtype=sample_array.dtype, + device=_array_device_or_none(sample_array), + ) + embed = self.embedding_net(backend_econf_tebd) if self.padding: embed_pad = xp.zeros( (1, embed.shape[-1]), diff --git a/source/tests/consistent/test_type_embedding.py b/source/tests/consistent/test_type_embedding.py index 08ebd36040..39e196006e 100644 --- a/source/tests/consistent/test_type_embedding.py +++ b/source/tests/consistent/test_type_embedding.py @@ -242,3 +242,53 @@ def atol(self) -> float: return 1e-1 else: raise ValueError(f"Unknown precision: {precision}") + + +@unittest.skipUnless(INSTALLED_PT, "PyTorch is not installed") +class TestDpmodelElectronicConfigurationBackend(unittest.TestCase): + """Exercise dpmodel with backend weights but portable NumPy constants.""" + + def test_econf_input_follows_torch_weights(self) -> None: + type_embedding = TypeEmbedNetDP( + ntypes=2, + neuron=[4, 4], + precision="float64", + use_econf_tebd=True, + type_map=["O", "H"], + seed=20260717, + ) + self.assertIsInstance(type_embedding.econf_tebd, np.ndarray) + numpy_reference = TypeEmbedNetDP.deserialize(type_embedding.serialize()) + expected = numpy_reference() + + # Model converters can replace the trainable arrays without touching + # portable constants. This reproduces that mixed-backend boundary + # directly, without relying on a wrapper that eagerly converts both. + for layer in type_embedding.embedding_net.layers: + for name in ("w", "b", "idt"): + value = getattr(layer, name) + if isinstance(value, np.ndarray): + setattr(layer, name, torch.as_tensor(value)) + + sample_weight = type_embedding.embedding_net[0]["w"] + result = type_embedding() + + self.assertIsInstance(result, torch.Tensor) + self.assertEqual(result.dtype, sample_weight.dtype) + self.assertEqual(result.device, sample_weight.device) + np.testing.assert_allclose(result.detach().cpu().numpy(), expected) + + # change_type_map intentionally regenerates the portable NumPy table; + # call-time conversion must continue to follow the existing weights. + new_type_map = ["C", "H", "O"] + numpy_reference.change_type_map(new_type_map) + type_embedding.change_type_map(new_type_map) + remapped_result = type_embedding() + + self.assertIsInstance(type_embedding.econf_tebd, np.ndarray) + self.assertIsInstance(remapped_result, torch.Tensor) + self.assertEqual(remapped_result.dtype, sample_weight.dtype) + self.assertEqual(remapped_result.device, sample_weight.device) + np.testing.assert_allclose( + remapped_result.detach().cpu().numpy(), numpy_reference() + )