Support JAX multitask training and finetuning

.github/workflows/build_wheel.yml

@@ -59,7 +59,7 @@ jobs:
         run: curl --proto '=https' --tlsv1.2 -LsSf https://github.com/astral-sh/uv/releases/download/0.2.24/uv-installer.sh | sh
         if: runner.os != 'Linux'
       - name: Build wheels
-        uses: pypa/cibuildwheel@v3.3
+        uses: pypa/cibuildwheel@v3.4
         env:
           CIBW_BUILD_VERBOSITY: 1
           CIBW_ARCHS: all

.github/workflows/package_c.yml

@@ -21,8 +21,8 @@ jobs:
     strategy:
       matrix:
         include:
-          - tensorflow_build_version: "2.18"
-            tensorflow_version: ""
+          - tensorflow_build_version: "2.20"
+            tensorflow_version: "==2.20.*"
             filename: libdeepmd_c.tar.gz
     steps:
       - name: Free Disk Space (Ubuntu)

.pre-commit-config.yaml

@@ -39,14 +39,6 @@ repos:
       - id: ruff-format
         exclude: ^source/3rdparty
         types_or: [python, pyi, jupyter]
-  - repo: https://github.com/pycqa/flake8
-    # flake8 cannot autofix
-    rev: "7.3.0"
-    hooks:
-      - id: flake8
-        additional_dependencies:
-          - torchfix==0.7.0
-          - flake8-pyproject==1.2.3
   # numpydoc
   - repo: https://github.com/Carreau/velin
     rev: 0.0.12

deepmd/dpmodel/array_api.py

@@ -15,6 +15,21 @@
 Array = np.ndarray | Any  # Any to support JAX, PyTorch, etc. arrays
 
 
+def _xp_arange(
+    xp: Any,
+    size: int,
+    dtype: Any,
+    like: Array,
+) -> Array:
+    if array_api_compat.is_jax_array(like):
+        return xp.arange(size, dtype=dtype)  # pylint: disable=no-explicit-device
+    return xp.arange(
+        size,
+        dtype=dtype,
+        device=array_api_compat.device(like),
+    )
+
+
 # array api adds take_along_axis in https://github.com/data-apis/array-api/pull/816
 # but it hasn't been released yet
 # below is a pure Python implementation of take_along_axis
@@ -51,8 +66,7 @@ def xp_take_along_axis(arr: Array, indices: Array, axis: int) -> Array:
     else:
         indices = xp.reshape(indices, (0, 0))
 
-    dev = array_api_compat.device(indices)
-    offset = (xp.arange(indices.shape[0], dtype=indices.dtype, device=dev) * m)[
+    offset = (_xp_arange(xp, indices.shape[0], indices.dtype, indices) * m)[
         :, xp.newaxis
     ]
     indices = xp.reshape(offset + indices, (-1,))
@@ -78,7 +92,7 @@ def xp_scatter_sum(input: Array, dim: int, index: Array, src: Array) -> Array:
     xp = array_api_compat.array_namespace(input)
 
     # Create flat index array matching input shape
-    idx = xp.arange(input.size, dtype=xp.int64, device=array_api_compat.device(input))
+    idx = _xp_arange(xp, input.size, xp.int64, input)
     idx = xp.reshape(idx, input.shape)
 
     # Get flat indices where we want to add values

deepmd/dpmodel/descriptor/dpa3.py

@@ -487,8 +487,14 @@ def change_type_map(
             )
         repflow.ntypes = self.ntypes
         repflow.reinit_exclude(self.exclude_types)
-        repflow["davg"] = repflow["davg"][remap_index]
-        repflow["dstd"] = repflow["dstd"][remap_index]
+        xp = array_api_compat.array_namespace(repflow["davg"])
+        remap_index_array = xp.asarray(
+            remap_index,
+            dtype=xp.int32,
+            device=array_api_compat.device(repflow["davg"]),
+        )
+        repflow["davg"] = repflow["davg"][remap_index_array]
+        repflow["dstd"] = repflow["dstd"][remap_index_array]
 
     @property
     def dim_out(self) -> int:

deepmd/dpmodel/descriptor/repflows.py

@@ -58,6 +58,28 @@
 )
 
 
+def _maybe_apply_jax_placeholder_sharding(array: Array) -> Array:
+    """Keep static placeholder tensors replicated when a mesh context exists.
+
+    JAX training runs under an explicit mesh and benefits from pinning these
+    dummy index tensors to replicated sharding. Export/freeze does not create
+    a mesh context, so the same sharding constraint must be skipped there.
+    """
+    from deepmd.jax.env import (
+        jax,
+    )
+
+    try:
+        return jax.lax.with_sharding_constraint(
+            array, jax.sharding.PartitionSpec(None, None)
+        )
+    except RuntimeError as exc:
+        message = str(exc)
+        if "non-empty mesh" not in message:
+            raise
+        return array
+
+
 @DescriptorBlock.register("se_repflow")
 class DescrptBlockRepflows(NativeOP, DescriptorBlock):
     r"""
@@ -595,12 +617,25 @@ def call(
             # n_angle x 1
             a_sw = (a_sw[:, :, :, None] * a_sw[:, :, None, :])[a_nlist_mask]
         else:
+            dummy_edge_count = max(1, self.nnei)
+            dummy_angle_count = max(1, self.a_sel * self.a_sel)
             edge_index = xp.zeros(
-                [2, 1], dtype=nlist.dtype, device=array_api_compat.device(nlist)
+                [2, dummy_edge_count],
+                dtype=nlist.dtype,
+                device=array_api_compat.device(nlist),
             )
             angle_index = xp.zeros(
-                [3, 1], dtype=nlist.dtype, device=array_api_compat.device(nlist)
+                [3, dummy_angle_count],
+                dtype=nlist.dtype,
+                device=array_api_compat.device(nlist),
             )
+            if array_api_compat.is_jax_namespace(xp):
+                # These are placeholders in the static-selection path and should
+                # stay replicated instead of inheriting the active natoms mesh.
+                # Freeze/export does not run under a mesh, so skip the
+                # constraint there and keep the placeholders as-is.
+                edge_index = _maybe_apply_jax_placeholder_sharding(edge_index)
+                angle_index = _maybe_apply_jax_placeholder_sharding(angle_index)
 
         # get edge and angle embedding
         # nb x nloc x nnei x e_dim [OR] n_edge x e_dim

deepmd/dpmodel/descriptor/se_e2_a.py

@@ -461,38 +461,21 @@ def call(
                 gg = self.cal_g(ss, (tt,))
                 gr += xp.sum(gg[:, :, :, None] * tr[:, :, None, :], axis=1)
         else:
-            # Sort atoms by center type so each type forms a contiguous block.
-            # Slice indexing (arr[s:e]) is array-api compatible and lets us
-            # run cal_g only on atoms of the matching center type, keeping the
-            # same O(nf*nloc) total embedding cost as the original numpy code.
             atype_loc = xp.reshape(atype_ext[:, :nloc], (nf * nloc,))
-            sort_idx = xp.argsort(atype_loc)
-            unsort_idx = xp.argsort(sort_idx)
-            rr_s = xp.take(rr, sort_idx, axis=0)
-            mask_s = xp.take(exclude_mask, sort_idx, axis=0)
-            dev = array_api_compat.device(coord_ext)
-            gr_s = xp.zeros([nf * nloc, ng, 4], dtype=input_dtype, device=dev)
-            # Per-type boundaries in sorted order
-            type_ends = []
-            offset = 0
             for ti in range(self.ntypes):
-                offset += int(xp.sum(xp.astype(atype_loc == ti, xp.int32)))
-                type_ends.append(offset)
-            type_starts = [0, *type_ends[:-1]]
-            for ti in range(self.ntypes):
-                s, e = type_starts[ti], type_ends[ti]
-                if s == e:
-                    continue
+                center_mask = xp.astype(atype_loc == ti, input_dtype)
+                center_mask = xp.reshape(center_mask, (nf * nloc, 1, 1))
                 for tt in range(self.ntypes):
-                    mm = mask_s[s:e, sec[tt] : sec[tt + 1]]
-                    tr = rr_s[s:e, sec[tt] : sec[tt + 1], :]
+                    mm = exclude_mask[:, sec[tt] : sec[tt + 1]]
+                    tr = rr[:, sec[tt] : sec[tt + 1], :]
                     tr = tr * xp.astype(mm[:, :, None], tr.dtype)
                     ss = tr[..., 0:1]
                     gg = self.cal_g(ss, (ti, tt))
-                    gr_s[s:e] = gr_s[s:e] + xp.sum(
-                        gg[:, :, :, None] * tr[:, :, None, :], axis=1
+                    gr = (
+                        gr
+                        + xp.sum(gg[:, :, :, None] * tr[:, :, None, :], axis=1)
+                        * center_mask
                     )
-            gr = xp.take(gr_s, unsort_idx, axis=0)
         gr = xp.reshape(gr, (nf, nloc, ng, 4))
         # nf x nloc x ng x 4
         gr /= self.nnei

deepmd/dpmodel/fitting/general_fitting.py

@@ -368,7 +368,13 @@ def change_type_map(
             self.bias_atom_e = np.concatenate(
                 [self.bias_atom_e, extend_bias_atom_e], axis=0
             )
-        self.bias_atom_e = self.bias_atom_e[remap_index]
+        xp = array_api_compat.array_namespace(self.bias_atom_e)
+        remap_index_array = xp.asarray(
+            remap_index,
+            dtype=xp.int32,
+            device=array_api_compat.device(self.bias_atom_e),
+        )
+        self.bias_atom_e = self.bias_atom_e[remap_index_array]
 
     def __setitem__(self, key: str, value: Any) -> None:
         if key in ["bias_atom_e"]:

deepmd/dpmodel/utils/nlist.py

@@ -119,14 +119,15 @@ def build_neighbor_list(
     device = array_api_compat.device(diff)
     if array_api_compat.is_jax_namespace(xp):
         # fix jax sharding "list index out of range"
-        from jax.sharding import PartitionSpec as P, NamedSharding
+        from jax.sharding import (
+            NamedSharding,
+        )
+        from jax.sharding import PartitionSpec as P
 
         if isinstance(device, NamedSharding):
             device = NamedSharding(device.mesh, P())
     # if central atom has two zero distances, sorting sometimes can not exclude itself
-    rr -= xp.eye(nloc, nall, dtype=diff.dtype, device=device)[
-        xp.newaxis, :, :
-    ]
+    rr -= xp.eye(nloc, nall, dtype=diff.dtype, device=device)[xp.newaxis, :, :]
     nlist = xp.argsort(rr, axis=-1)
     rr = xp.sort(rr, axis=-1)
     rr = rr[:, :, 1:]
@@ -309,16 +310,20 @@ def extend_coord_with_ghosts(
     if array_api_compat.is_jax_namespace(xp):
         # fix jax: Sharding is only valid for values of rank at least 2,
         # but was applied to a value of rank 1.
-        from jax.sharding import PartitionSpec as P, NamedSharding
+        from jax.sharding import (
+            NamedSharding,
+        )
+        from jax.sharding import PartitionSpec as P
 
         if isinstance(device, NamedSharding):
-            device_nloc = NamedSharding(device.mesh, P(device.spec[1]))
+            if len(device.spec) > 1:
+                device_nloc = NamedSharding(device.mesh, P(device.spec[1]))
+            else:
+                device_nloc = NamedSharding(device.mesh, P())
             device_none = NamedSharding(device.mesh, P())
     # int64 for index
     aidx = xp.tile(
-        xp.arange(nloc, dtype=xp.int64, device=device_nloc)[
-            xp.newaxis, :
-        ],
+        xp.arange(nloc, dtype=xp.int64, device=device_nloc)[xp.newaxis, :],
         (nf, 1),
     )
     if cell is None:

deepmd/dpmodel/utils/region.py

@@ -77,13 +77,14 @@ def normalize_coord(
     device = array_api_compat.device(icoord)
     if array_api_compat.is_jax_namespace(xp):
         # fix jax sharding "list index out of range"
-        from jax.sharding import PartitionSpec as P, NamedSharding
+        from jax.sharding import (
+            NamedSharding,
+        )
+        from jax.sharding import PartitionSpec as P
 
         if isinstance(device, NamedSharding):
             device = NamedSharding(device.mesh, P())
-    icoord = xp.remainder(
-        icoord, xp.ones((), dtype=icoord.dtype, device=device)
-    )
+    icoord = xp.remainder(icoord, xp.ones((), dtype=icoord.dtype, device=device))
     return inter2phys(icoord, cell)
 
 

deepmd/dpmodel/utils/type_embed.py

@@ -226,7 +226,12 @@ def change_type_map(
                     [first_layer_matrix, extend_type_params], axis=0
                 )
 
-            first_layer_matrix = first_layer_matrix[remap_index]
+            remap_index_array = xp.asarray(
+                remap_index,
+                dtype=xp.int32,
+                device=array_api_compat.device(first_layer_matrix),
+            )
+            first_layer_matrix = first_layer_matrix[remap_index_array]
             new_ntypes = len(type_map)
             eye_vector = xp.eye(
                 new_ntypes,

deepmd/jax/descriptor/hybrid.py

@@ -3,11 +3,15 @@
     Any,
 )
 
+import array_api_compat
 from packaging.version import (
     Version,
 )
 
 from deepmd.dpmodel.descriptor.hybrid import DescrptHybrid as DescrptHybridDP
+from deepmd.dpmodel.utils.nlist import (
+    nlist_distinguish_types,
+)
 from deepmd.jax.common import (
     ArrayAPIVariable,
     flax_module,
@@ -36,3 +40,44 @@ def __setattr__(self, name: str, value: Any) -> None:
                 value = nnx.List([nnx.data(item) for item in value])
 
         return super().__setattr__(name, value)
+
+    def call(
+        self, *args: Any, **kwargs: Any
+    ) -> tuple[Any, Any | None, Any | None, Any | None, Any | None]:
+        if len(args) < 3:
+            return super().call(*args, **kwargs)
+        if len(args) > 4:
+            return super().call(*args, **kwargs)
+        if kwargs and set(kwargs) != {"mapping"}:
+            return super().call(*args, **kwargs)
+        coord_ext, atype_ext, nlist = args[:3]
+        mapping = kwargs.pop("mapping", args[3] if len(args) == 4 else None)
+        xp = array_api_compat.array_namespace(coord_ext, atype_ext, nlist)
+        out_descriptor = []
+        out_gr = []
+        out_g2 = None
+        out_h2 = None
+        out_sw = None
+        if self.sel_no_mixed_types is not None:
+            nl_distinguish_types = nlist_distinguish_types(
+                nlist,
+                atype_ext,
+                self.sel_no_mixed_types,
+            )
+        else:
+            nl_distinguish_types = None
+        for descrpt, nci in zip(self.descrpt_list, self.nlist_cut_idx, strict=True):
+            nci_value = getattr(nci, "value", nci)
+            if self.mixed_types() == descrpt.mixed_types():
+                nl = xp.take(nlist, nci_value, axis=2)
+            else:
+                assert nl_distinguish_types is not None
+                nl = nl_distinguish_types[:, :, nci_value]
+            odescriptor, gr, g2, h2, sw = descrpt(coord_ext, atype_ext, nl, mapping)
+            out_descriptor.append(odescriptor)
+            if gr is not None:
+                out_gr.append(gr)
+
+        out_descriptor = xp.concat(out_descriptor, axis=-1)
+        out_gr = xp.concat(out_gr, axis=-2) if out_gr else None
+        return out_descriptor, out_gr, out_g2, out_h2, out_sw

deepmd/jax/entrypoints/freeze.py

@@ -8,6 +8,7 @@
 
 from deepmd.jax.utils.serialization import (
     deserialize_to_file,
+    select_model_branch,
     serialize_from_file,
 )
 
@@ -16,6 +17,8 @@ def freeze(
     *,
     checkpoint_folder: str,
     output: str,
+    head: str | None = None,
+    model_branch: str | None = None,
     hessian: bool = False,
     **kwargs: Any,
 ) -> None:
@@ -37,6 +40,9 @@ def freeze(
         checkpoint_folder = checkpoint_meta.read_text().strip()
     if Path(checkpoint_folder).is_dir():
         data = serialize_from_file(checkpoint_folder)
+        selected_branch = model_branch or head
+        if selected_branch and "model_dict" in data["model_def_script"]:
+            data = select_model_branch(data, selected_branch)
         deserialize_to_file(output, data, hessian=hessian)
     else:
         raise FileNotFoundError(f"Checkpoint {checkpoint_folder} does not exist.")

deepmd/jax/entrypoints/main.py

@@ -5,10 +5,6 @@
 from pathlib import (
     Path,
 )
-from typing import (
-    Optional,
-    Union,
-)
 
 from deepmd.backend.suffix import (
     format_model_suffix,
@@ -29,7 +25,7 @@
 __all__ = ["main"]
 
 
-def main(args: Optional[Union[list[str], argparse.Namespace]] = None) -> None:
+def main(args: list[str] | argparse.Namespace | None = None) -> None:
     """DeePMD-Kit entry point.
 
     Parameters