fix(jax): avoid device lookup during type embedding tracing (#5736)

## Summary

Fix a JAX/Flax NNX tracing failure in `TypeEmbedNet.call()` by avoiding
a hard requirement that traced values expose a readable `.device`
attribute.

## Background

JAX training can call `TypeEmbedNet.call()` from inside `nnx.jit` /
`nnx.grad`. In that context the type-embedding weights are represented
by NNX/JAX traced values such as `DynamicJaxprTracer`. The previous code
passed

```python
device=array_api_compat.device(sample_array)
```

to `xp.eye`, and did the same for the padding `xp.zeros`. For an NNX
traced parameter, `array_api_compat.device(...)` eventually tries to
read `.device`; the tracer does not provide that attribute during
tracing, so the training step fails with an error like:

```text
AttributeError: DynamicJaxprTracer has no attribute device
```

or, through JAX core wrapping:

```text
AttributeError: 'ShapedArray' object has no attribute 'device'
```

## Change

Add a small local helper in `deepmd/dpmodel/utils/type_embed.py` that
returns `array_api_compat.device(array)` when available, and falls back
to `None` when the backend value has no readable device during tracing.
Passing `device=None` lets JAX create the constants on its default
traced device, while preserving the existing explicit-device behavior
for backends that expose it normally.

This keeps the change limited to the type-embedding constants that
triggered the crash, rather than changing global array API device
handling.

## Tests

Added `source/tests/jax/test_type_embed.py` to cover both:

- `TypeEmbedNet.call()` under `nnx.jit`
- `TypeEmbedNet.call()` under `nnx.jit` + `nnx.grad`

The old implementation fails this test during tracing before producing
an output.

Validation run locally:

```bash
pytest source/tests/jax/test_type_embed.py -v
ruff check .
ruff format .
```


<!-- This is an auto-generated comment: release notes by coderabbit.ai
-->

## Summary by CodeRabbit

* **Bug Fixes**
* Improved compatibility for model execution on array backends that do
not support device detection in all cases.
* Reduced failures when padding or initializing arrays during type
embedding operations.

* **Tests**
* Added JAX coverage to verify the type embedding model runs under JIT,
supports gradient tracing, returns the expected output shape, and
produces valid numeric results.

<!-- end of auto-generated comment: release notes by coderabbit.ai -->

Author: njzjz

deepmd/dpmodel/utils/type_embed.py

@@ -24,6 +24,13 @@
 )
 
 
+def _array_device_or_none(array: Array) -> Any:
+    try:
+        return array_api_compat.device(array)
+    except AttributeError:
+        return None
+
+
 class TypeEmbedNet(NativeOP):
     r"""Type embedding network.
 
@@ -104,7 +111,7 @@ def call(self) -> Array:
                 xp.eye(
                     self.ntypes,
                     dtype=sample_array.dtype,
-                    device=array_api_compat.device(sample_array),
+                    device=_array_device_or_none(sample_array),
                 )
             )
         else:
@@ -113,7 +120,7 @@ def call(self) -> Array:
             embed_pad = xp.zeros(
                 (1, embed.shape[-1]),
                 dtype=embed.dtype,
-                device=array_api_compat.device(embed),
+                device=_array_device_or_none(embed),
             )
             embed = xp.concat([embed, embed_pad], axis=0)
         return embed

source/tests/jax/test_type_embed.py

@@ -0,0 +1,44 @@
+# SPDX-License-Identifier: LGPL-3.0-or-later
+import unittest
+
+import numpy as np
+
+from deepmd.jax.env import (
+    jnp,
+    nnx,
+)
+from deepmd.jax.utils.type_embed import (
+    TypeEmbedNet,
+)
+
+
+class TestJAXTypeEmbedNet(unittest.TestCase):
+    def test_call_supports_nnx_jit_and_grad_tracing(self) -> None:
+        type_embedding = TypeEmbedNet(
+            ntypes=2,
+            neuron=[4],
+            padding=True,
+            activation_function="Linear",
+            precision="float32",
+            seed=1,
+        )
+
+        @nnx.jit
+        def forward(model):
+            return model.call()
+
+        @nnx.jit
+        def grad(model):
+            def loss(model):
+                return jnp.sum(model.call())
+
+            return nnx.grad(loss)(model)
+
+        out = forward(type_embedding)
+        self.assertEqual(tuple(out.shape), (3, 4))
+        self.assertIsNotNone(grad(type_embedding))
+        self.assertFalse(np.any(np.isnan(np.asarray(out))))
+
+
+if __name__ == "__main__":
+    unittest.main()