feat(pt_expt): call-time DeepEval auto ladder with nf==1 vesin gate

Extract resolve_auto_graph_builder for inference and select vesin only
for single-frame batches, matching _select_neighbor_builder. Multi-frame
auto stays on nv/dense so auto_batch_size does not hit the per-frame loop.

Signed-off-by: shaurya2k06 <shaurya2k06@gmail.com>

Author: Shaurya2k06

.github/workflows/build_wheel.yml

@@ -56,7 +56,9 @@ jobs:
           # https://github.com/pypa/setuptools_scm/issues/480
           fetch-depth: 0
       - name: Install uv
-        run: curl --proto '=https' --tlsv1.2 -LsSf https://github.com/astral-sh/uv/releases/download/0.2.24/uv-installer.sh | sh
+        uses: astral-sh/setup-uv@c771a70e6277c0a99b617c7a806ffedaca235ff9 # v9.0.0
+        with:
+          version: "0.12.0"
         if: runner.os != 'Linux'
       - name: Build wheels
         uses: pypa/cibuildwheel@v4.1
@@ -68,6 +70,9 @@ jobs:
           CUDA_VERSION: ${{ matrix.cuda_version }}
           DP_PKG_NAME: ${{ matrix.dp_pkg_name }}
           CIBW_BUILD_FRONTEND: "build[uv]"
+          # uv defaults to 50 concurrent downloads, which can exhaust the
+          # macOS runner's socket buffers while installing wheel test extras.
+          UV_CONCURRENT_DOWNLOADS: ${{ startsWith(matrix.platform_id, 'macosx_') && '8' || '50' }}
       - uses: actions/upload-artifact@v7
         with:
           name: cibw-cp${{ matrix.python }}-${{ matrix.platform_id }}-cu${{ matrix.cuda_version }}-${{ strategy.job-index }}

.github/workflows/test_cc.yml

@@ -15,7 +15,12 @@ jobs:
   testcc:
     name: Test C++
     runs-on: ${{ matrix.check_memleak == false && github.repository_owner == 'deepmodeling' && 'cpu' || 'ubuntu-22.04' }}
+    # Paddle's external interface download is occasionally unavailable. Keep
+    # exercising these configurations without making that outage blocking.
+    continue-on-error: ${{ matrix.enable_paddle }}
     strategy:
+      # An allowed Paddle failure must not cancel the other backend jobs.
+      fail-fast: false
       matrix:
         # Only regular jobs run the LAMMPS tests. Use self-hosted CPU runners
         # for those jobs in deepmodeling, with a hosted fallback for forks.

.github/workflows/test_cuda.yml

@@ -13,6 +13,11 @@ concurrency:
   group: ${{ github.workflow }}-${{ github.ref || github.run_id }}
   cancel-in-progress: true
 name: Test CUDA
+env:
+  # Temporarily disable Paddle in CUDA CI while its external downloads are
+  # unreliable. Keep this switch explicit so the coverage can be restored.
+  DP_CI_ALLOW_MISSING_PADDLE: "1"
+  DP_ENABLE_PADDLE: "0"
 jobs:
   test_python:
     name: Test Python on CUDA
@@ -53,7 +58,6 @@ jobs:
       - run: |
           export PYTORCH_ROOT=$(python -c 'import torch;print(torch.__path__[0])')
           export TENSORFLOW_ROOT=$(python -c 'import importlib.util,pathlib;print(pathlib.Path(importlib.util.find_spec("tensorflow").origin).parent)')
-          pip install --find-links "https://www.paddlepaddle.org.cn/packages/nightly/cu126/paddlepaddle-gpu/" --index-url https://pypi.org/simple --trusted-host www.paddlepaddle.org.cn --trusted-host paddlepaddle.org.cn "paddlepaddle-gpu==3.4.0.dev20260310"
           source/install/uv_with_retry.sh pip install --system -v -e .[gpu,test,lmp,cu12,torch,jax] mpi4py --reinstall-package deepmd-kit
           # See https://github.com/jax-ml/jax/issues/29042
           source/install/uv_with_retry.sh pip install --system -U 'nvidia-cublas-cu12>=12.9.0.13'
@@ -62,7 +66,7 @@ jobs:
           DP_ENABLE_NATIVE_OPTIMIZATION: 1
           DP_ENABLE_PYTORCH: 1
       - run: dp --version
-      - run: python -m pytest source/tests
+      - run: python -m pytest source/tests --ignore=source/tests/pd
         env:
           NUM_WORKERS: 0
           CUDA_VISIBLE_DEVICES: 0
@@ -106,7 +110,6 @@ jobs:
       - run: |
           export PYTORCH_ROOT=$(python -c 'import torch;print(torch.__path__[0])')
           export TENSORFLOW_ROOT=$(python -c 'import importlib.util,pathlib;print(pathlib.Path(importlib.util.find_spec("tensorflow").origin).parent)')
-          pip install --find-links "https://www.paddlepaddle.org.cn/packages/nightly/cu126/paddlepaddle-gpu/" --index-url https://pypi.org/simple --trusted-host www.paddlepaddle.org.cn --trusted-host paddlepaddle.org.cn "paddlepaddle-gpu==3.4.0.dev20260310"
           source/install/uv_with_retry.sh pip install --system -v -e .[gpu,test,lmp,cu12,torch,jax] mpi4py --reinstall-package deepmd-kit
           # See https://github.com/jax-ml/jax/issues/29042
           source/install/uv_with_retry.sh pip install --system -U 'nvidia-cublas-cu12>=12.9.0.13'
@@ -127,12 +130,10 @@ jobs:
           CMAKE_GENERATOR: Ninja
           DP_VARIANT: cuda
           DP_USE_MPICH2: 1
+          ENABLE_PADDLE: FALSE
       - run: |
           export LD_LIBRARY_PATH=$CUDA_PATH/lib64:/usr/lib/x86_64-linux-gnu/:$GITHUB_WORKSPACE/dp_test/lib:$LD_LIBRARY_PATH
           export PATH=$GITHUB_WORKSPACE/dp_test/bin:$PATH
-          cp $GITHUB_WORKSPACE/source/build_tests/paddle_inference_install_dir/paddle/lib/* $GITHUB_WORKSPACE/dp_test/lib/
-          cp $GITHUB_WORKSPACE/source/build_tests/paddle_inference_install_dir/third_party/install/onednn/lib/* $GITHUB_WORKSPACE/dp_test/lib/
-          cp $GITHUB_WORKSPACE/source/build_tests/paddle_inference_install_dir/third_party/install/mklml/lib/* $GITHUB_WORKSPACE/dp_test/lib/
           python -m pytest -s source/lmp/tests || (cat log.lammps && exit 1)
           python -m pytest source/ipi/tests
         env:
@@ -141,6 +142,7 @@ jobs:
           TF_INTER_OP_PARALLELISM_THREADS: 1
           LAMMPS_PLUGIN_PATH: ${{ github.workspace }}/dp_test/lib/deepmd_lmp
           CUDA_VISIBLE_DEVICES: 0
+          ENABLE_PADDLE: 0
   pass:
     name: Pass testing on CUDA
     needs: [test_python, test_cc]

deepmd/dpmodel/array_api.py

@@ -201,6 +201,22 @@ def xp_add_at(x: Array, indices: Array, values: Array) -> Array:
         import torch
 
         return torch.index_add(x, 0, indices, values)
+    elif getattr(xp, "__name__", "") == "deepmd._vendors.ndtensorflow":
+        import tensorflow as tf
+
+        x_tensor = x.unwrap()
+        indices_tensor = tf.reshape(tf.cast(indices.unwrap(), tf.int64), (-1,))
+        values_tensor = values.unwrap()
+        # unsorted_segment_sum rather than scatter_nd: both accumulate repeated
+        # indices, but scatter_nd rejects a destination with no elements even
+        # when the updates are empty too, which a descriptor call with zero
+        # edges legitimately produces.
+        updates = tf.math.unsorted_segment_sum(
+            values_tensor,
+            indices_tensor,
+            tf.shape(x_tensor, out_type=tf.int64)[0],
+        )
+        return xp.asarray(x_tensor + updates)
     else:
         # Fallback for array_api_strict: use basic indexing only
         # may need a more efficient way to do this
@@ -270,6 +286,52 @@ def xp_maximum_at(x: Array, indices: Array, values: Array) -> Array:
         return torch.scatter_reduce(
             x, 0, index, values, reduce="amax", include_self=True
         )
+    elif getattr(xp, "__name__", "") == "deepmd._vendors.ndtensorflow":
+        import tensorflow as tf
+
+        x_tensor = x.unwrap()
+        indices_tensor = tf.reshape(tf.cast(indices.unwrap(), tf.int64), (-1,))
+        values_tensor = values.unwrap()
+        reduced = tf.math.unsorted_segment_max(
+            values_tensor,
+            indices_tensor,
+            tf.shape(x_tensor, out_type=tf.int64)[0],
+        )
+        if values_tensor.dtype.is_floating:
+            # TensorFlow uses the lowest finite value as the identity of
+            # unsorted_segment_max. Restore the true maximum-at identity when
+            # every update for a touched segment element is negative infinity.
+            all_negative_infinity = (
+                tf.math.unsorted_segment_min(
+                    tf.cast(
+                        tf.math.is_inf(values_tensor) & (values_tensor < 0),
+                        tf.int32,
+                    ),
+                    indices_tensor,
+                    tf.shape(x_tensor, out_type=tf.int64)[0],
+                )
+                > 0
+            )
+            reduced = tf.where(
+                all_negative_infinity,
+                tf.cast(float("-inf"), values_tensor.dtype),
+                reduced,
+            )
+        segment_counts = tf.math.unsorted_segment_sum(
+            tf.ones_like(indices_tensor, dtype=tf.int32),
+            indices_tensor,
+            tf.shape(x_tensor, out_type=tf.int64)[0],
+        )
+        touched = segment_counts > 0
+        touched_shape = tf.concat(
+            [
+                tf.reshape(tf.shape(x_tensor, out_type=tf.int64)[0], (1,)),
+                tf.ones(tf.rank(x_tensor) - 1, dtype=tf.int64),
+            ],
+            axis=0,
+        )
+        touched = tf.reshape(touched, touched_shape)
+        return xp.asarray(tf.where(touched, tf.maximum(x_tensor, reduced), x_tensor))
     else:
         # Fallback for array_api_strict: basic indexing only.
         n = indices.shape[0]
@@ -337,12 +399,12 @@ def xp_setitem_at(x: Array, mask: Array, values: Array) -> Array:
 def xp_uniform(like: Array, size: int, low: float = 0.0, high: float = 1.0) -> Array:
     """Draw ``size`` uniform samples in ``[low, high)`` on ``like``'s device.
 
-    Each backend uses its own generator: torch draws with ``torch.rand`` (as
-    pt does, so ``setup_seed`` replays it, with no host copy -- and a host
-    draw would freeze to a constant under tracing); other backends use
-    :mod:`deepmd.utils.random`, which ``setup_seed`` also seeds. Draws are
-    therefore not comparable across backends -- use only for a per-forward
-    random stream, never where a parity test looks.
+    Each backend uses its own generator: TensorFlow draws with
+    ``tf.random.uniform`` so traced graphs advance the runtime RNG, torch draws
+    with ``torch.rand`` (so ``setup_seed`` replays it without a host copy), and
+    other backends use :mod:`deepmd.utils.random`. Draws are therefore not
+    comparable across backends -- use only for a per-forward random stream,
+    never where a parity test looks.
 
     Parameters
     ----------
@@ -360,6 +422,18 @@ def xp_uniform(like: Array, size: int, low: float = 0.0, high: float = 1.0) -> A
     Array
         Samples of shape ``(size,)`` matching ``like``.
     """
+    xp = array_api_compat.array_namespace(like)
+    if getattr(xp, "__name__", "") == "deepmd._vendors.ndtensorflow":
+        import tensorflow as tf
+
+        sample_shape = tf.reshape(tf.cast(size, tf.int32), (1,))
+        samples = tf.random.uniform(
+            sample_shape,
+            minval=low,
+            maxval=high,
+            dtype=like.dtype,
+        )
+        return xp.asarray(samples)
     if array_api_compat.is_torch_array(like):
         import torch
 
@@ -368,7 +442,6 @@ def xp_uniform(like: Array, size: int, low: float = 0.0, high: float = 1.0) -> A
         )
     from deepmd.utils import random as dp_random
 
-    xp = array_api_compat.array_namespace(like)
     drawn = np.asarray(dp_random.random(size)) * (high - low) + low
     return xp.astype(
         xp_asarray_nodetach(xp, drawn, device=array_api_compat.device(like)),

deepmd/dpmodel/atomic_model/base_atomic_model.py

@@ -190,6 +190,54 @@ def supports_native_spin(self) -> bool:
         """
         return False
 
+    def has_message_passing_across_ranks(self) -> bool:
+        """Whether multi-rank inference needs a cross-rank ghost exchange.
+
+        Generic capability (concrete default ``False``): the export layer
+        consults it instead of reaching for a descriptor, so the answer
+        stays correct for descriptor-less models and compositions.
+        """
+        return False
+
+    def supports_edge_parallel(self) -> bool:
+        """Whether this atomic model can run under MPI domain decomposition.
+
+        Default ``True``; a model folding state no single rank observes
+        (e.g. SFPG bridging before its exchange lands) overrides to False.
+        """
+        return True
+
+    def dense_lower_supports_comm(self) -> bool:
+        """Whether the DENSE (nlist) lower implements comm_dict exchange.
+
+        Default ``True`` — dense comm is the production multi-rank path for
+        dpa2/dpa3; DPA4's dense adapter raises on comm_dict and overrides
+        via its descriptor.
+        """
+        return True
+
+    def uses_compact_edge_pairs(self) -> bool:
+        """Whether the graph lower emits compact ``center_edge_pairs``
+        (drives the torch>=2.6 unbacked-SymInt export guard).
+        """
+        return False
+
+    def graph_edge_dtype(self) -> str:
+        """Edge-geometry dtype the graph deployment artifact accepts.
+
+        ``"float64"`` is the model-agnostic ABI; geometrically compressed
+        float32 descriptors override to ``"float32"``.
+        """
+        return "float64"
+
+    def supports_graph_export(self) -> bool:
+        """Whether an exportable graph-lower implementation exists.
+
+        A compressed descriptor without its fused opaque operator cannot be
+        traced through the reference tabulation kernel.
+        """
+        return True
+
     def get_default_fparam(self) -> list[float] | None:
         """Get the default frame parameters."""
         return None
@@ -805,9 +853,22 @@ def _store_out_stat(
         self.out_std = out_std_data
 
     def _get_forward_wrapper_func(self) -> Callable[..., dict[str, np.ndarray]]:
-        """Get a forward wrapper of the atomic model for output bias calculation."""
+        """Get a forward wrapper of the atomic model for output bias calculation.
+
+        The wrapper starts from raw coordinates and therefore has to construct
+        the neighbor representation itself. It builds the one this model
+        declares through :meth:`uses_graph_lower`: a carry-all
+        ``NeighborGraph`` for graph-native models, whose neighbor count follows
+        the geometry, or the fixed-capacity neighbor list sized by
+        :meth:`get_sel` otherwise. Sizing a dense list from ``get_sel`` is not
+        merely wasteful for a graph-native model -- such a model reports no
+        finite capacity, so the allocation is unbounded.
+        """
         import array_api_compat
 
+        from deepmd.dpmodel.utils.neighbor_graph import (
+            build_neighbor_graph,
+        )
         from deepmd.dpmodel.utils.nlist import (
             extend_input_and_build_neighbor_list,
         )
@@ -841,31 +902,64 @@ def model_forward(
             if charge_spin is not None:
                 charge_spin = xp.asarray(charge_spin, device=device)
 
-            (
-                extended_coord,
-                extended_atype,
-                mapping,
-                nlist,
-            ) = extend_input_and_build_neighbor_list(
-                coord,
-                atype,
-                self.get_rcut(),
-                self.get_sel(),
-                mixed_types=self.mixed_types(),
-                box=box,
-                # exclusion is a nlist-BUILD transform (decision #18/A4);
-                # forward_common_atomic consumes a pre-excluded nlist.
-                pair_excl=self.pair_excl,
-            )
-            atomic_ret = self.forward_common_atomic(
-                extended_coord,
-                extended_atype,
-                nlist,
-                mapping=mapping,
-                fparam=fparam,
-                aparam=aparam,
-                charge_spin=charge_spin,
-            )
+            if self.uses_graph_lower():
+                nframes, nloc = atype.shape
+                # Pair exclusion is a neighbor-BUILD transform (decision
+                # #18/A4) on both routes; the graph builder folds it into
+                # ``edge_mask``.
+                graph = build_neighbor_graph(
+                    coord,
+                    atype,
+                    box,
+                    self.get_rcut(),
+                    pair_excl=self.pair_excl,
+                )
+                atomic_ret = self.forward_common_atomic_graph(
+                    graph,
+                    xp.reshape(atype, (-1,)),
+                    fparam=fparam,
+                    aparam=(
+                        xp.reshape(
+                            aparam,
+                            (nframes * nloc, self.get_dim_aparam()),
+                        )
+                        if aparam is not None
+                        else None
+                    ),
+                    charge_spin=charge_spin,
+                )
+                # The graph route works on a flat node axis; restore the
+                # per-frame layout the dense route returns.
+                atomic_ret = {
+                    kk: xp.reshape(vv, (nframes, nloc, *vv.shape[1:]))
+                    for kk, vv in atomic_ret.items()
+                }
+            else:
+                (
+                    extended_coord,
+                    extended_atype,
+                    mapping,
+                    nlist,
+                ) = extend_input_and_build_neighbor_list(
+                    coord,
+                    atype,
+                    self.get_rcut(),
+                    self.get_sel(),
+                    mixed_types=self.mixed_types(),
+                    box=box,
+                    # exclusion is a nlist-BUILD transform (decision #18/A4);
+                    # forward_common_atomic consumes a pre-excluded nlist.
+                    pair_excl=self.pair_excl,
+                )
+                atomic_ret = self.forward_common_atomic(
+                    extended_coord,
+                    extended_atype,
+                    nlist,
+                    mapping=mapping,
+                    fparam=fparam,
+                    aparam=aparam,
+                    charge_spin=charge_spin,
+                )
             # Convert outputs back to numpy arrays
             return {kk: to_numpy_array(vv) for kk, vv in atomic_ret.items()}