Add interface to multi-fitting architecture.

.github/workflows/build_wheel.yml

@@ -74,30 +74,27 @@ jobs:
         with:
           # https://github.com/pypa/setuptools_scm/issues/480
           fetch-depth: 0
+      - name: Install uv
+        run: curl -LsSf https://astral.sh/uv/install.sh | sh
+        if: runner.os != 'Linux'
       - uses: docker/setup-qemu-action@v3
         name: Setup QEMU
         if: matrix.platform_id == 'manylinux_aarch64' && matrix.os == 'ubuntu-latest'
       # detect version in advance. See #3168
-      - uses: actions/setup-python@v5
-        name: Install Python
-        with:
-          python-version: '3.11'
-          cache: 'pip'
-        if: matrix.dp_pkg_name == 'deepmd-kit-cu11'
       - run: |
-          python -m pip install setuptools_scm
-          python -c "from setuptools_scm import get_version;print('SETUPTOOLS_SCM_PRETEND_VERSION='+get_version())" >> $GITHUB_ENV
+          echo "SETUPTOOLS_SCM_PRETEND_VERSION=$(pipx run uv tool run --from setuptools_scm python -m setuptools_scm)" >> $GITHUB_ENV
           rm -rf .git
         if: matrix.dp_pkg_name == 'deepmd-kit-cu11'
       - name: Build wheels
-        uses: pypa/cibuildwheel@v2.18
+        uses: pypa/cibuildwheel@v2.19
         env:
           CIBW_BUILD_VERBOSITY: 1
           CIBW_ARCHS: all
           CIBW_BUILD: cp${{ matrix.python }}-${{ matrix.platform_id }}
           DP_VARIANT: ${{ matrix.dp_variant }}
           CUDA_VERSION: ${{ matrix.cuda_version }}
           DP_PKG_NAME: ${{ matrix.dp_pkg_name }}
+          CIBW_BUILD_FRONTEND: 'build[uv]'
       - uses: actions/upload-artifact@v4
         with:
           name: cibw-cp${{ matrix.python }}-${{ matrix.platform_id }}-cu${{ matrix.cuda_version }}-${{ strategy.job-index }}
@@ -109,14 +106,8 @@ jobs:
       - uses: actions/checkout@v4
         with:
           fetch-depth: 0
-      - uses: actions/setup-python@v5
-        name: Install Python
-        with:
-          python-version: '3.11'
-          cache: 'pip'
-      - run: python -m pip install build
       - name: Build sdist
-        run: python -m build --sdist
+        run: pipx run uv tool run --with build[uv] --from build python -m build --installer uv --sdist
 
       - uses: actions/upload-artifact@v4
         with:

.gitignore

@@ -44,3 +44,4 @@ build_cc_tests
 build_c_tests
 build_c/
 libdeepmd_c/
+.uv/

.pre-commit-config.yaml

@@ -29,7 +29,7 @@ repos:
         exclude: ^source/3rdparty
   - repo: https://github.com/astral-sh/ruff-pre-commit
     # Ruff version.
-    rev: v0.4.7
+    rev: v0.4.8
     hooks:
       - id: ruff
         args: ["--fix"]
@@ -52,7 +52,7 @@ repos:
       - id: blacken-docs
   # C++
   - repo: https://github.com/pre-commit/mirrors-clang-format
-    rev: v18.1.5
+    rev: v18.1.6
     hooks:
       - id: clang-format
         exclude: ^source/3rdparty|source/lib/src/gpu/cudart/.+\.inc

backend/find_tensorflow.py

@@ -156,11 +156,18 @@ def get_tf_requirement(tf_version: str = "") -> dict:
                 "tensorflow; platform_machine=='aarch64' or (platform_machine=='arm64' and platform_system == 'Darwin')",
                 # https://github.com/tensorflow/tensorflow/issues/61830
                 "tensorflow-cpu!=2.15.*; platform_system=='Windows'",
+                # TODO: build(wheel): unpin h5py on aarch64
+                # Revert after https://github.com/h5py/h5py/issues/2408 is fixed;
+                # or set UV_PREFER_BINARY when https://github.com/astral-sh/uv/issues/1794 is resolved.
+                # 3.6.0 is the first version to have aarch64 wheels.
+                "h5py>=3.6.0,<3.11.0; platform_system=='Linux' and platform_machine=='aarch64'",
                 *extra_requires,
             ],
             "gpu": [
                 "tensorflow",
                 "tensorflow-metal; platform_machine=='arm64' and platform_system == 'Darwin'",
+                # See above.
+                "h5py>=3.6.0,<3.11.0; platform_system=='Linux' and platform_machine=='aarch64'",
                 *extra_requires,
             ],
             **extra_select,

deepmd/dpmodel/atomic_model/make_base_atomic_model.py

@@ -7,6 +7,7 @@
     Dict,
     List,
     Optional,
+    Union,
 )
 
 from deepmd.dpmodel.output_def import (
@@ -87,7 +88,7 @@ def get_dim_aparam(self) -> int:
             """Get the number (dimension) of atomic parameters of this atomic model."""
 
         @abstractmethod
-        def get_sel_type(self) -> List[int]:
+        def get_sel_type(self) -> Union[List[int], List[List[int]]]:
             """Get the selected atom types of this model.
 
             Only atoms with selected atom types have atomic contribution

deepmd/pt/entrypoints/main.py

@@ -243,8 +243,10 @@ def train(FLAGS):
     if multi_task:
         config["model"], shared_links = preprocess_shared_params(config["model"])
 
+    multi_fitting_net = "fitting_net_dict" in config["model"]
+
     # argcheck
-    if not multi_task:
+    if not (multi_task or multi_fitting_net):
         config = update_deepmd_input(config, warning=True, dump="input_v2_compat.json")
         config = normalize(config)
 

deepmd/pt/model/atomic_model/__init__.py

@@ -26,6 +26,9 @@
 from .dp_atomic_model import (
     DPAtomicModel,
 )
+from .dp_multi_fitting_atomic_model import (
+    DPMultiFittingAtomicModel,
+)
 from .energy_atomic_model import (
     DPEnergyAtomicModel,
 )
@@ -50,4 +53,5 @@
     "DPPolarAtomicModel",
     "DPDipoleAtomicModel",
     "DPZBLLinearEnergyAtomicModel",
+    "DPMultiFittingAtomicModel",
 ]