feat: remove unused support_array_api decorator

Co-authored-by: njzjz <9496702+njzjz@users.noreply.github.com>

Author: Copilot

deepmd/dpmodel/array_api.py

@@ -1,9 +1,6 @@
 # SPDX-License-Identifier: LGPL-3.0-or-later
 """Utilities for the array API."""
 
-from collections.abc import (
-    Callable,
-)
 from typing import (
     Any,
 )
@@ -18,33 +15,6 @@
 Array = np.ndarray | Any  # Any to support JAX, PyTorch, etc. arrays
 
 
-def support_array_api(version: str) -> Callable:
-    """Mark a function as supporting the specific version of the array API.
-
-    Parameters
-    ----------
-    version : str
-        The version of the array API
-
-    Returns
-    -------
-    Callable
-        The decorated function
-
-    Examples
-    --------
-    >>> @support_array_api(version="2022.12")
-    ... def f(x):
-    ...     pass
-    """
-
-    def set_version(func: Callable) -> Callable:
-        func.array_api_version = version
-        return func
-
-    return set_version
-
-
 # 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

deepmd/dpmodel/utils/env_mat.py

@@ -10,15 +10,13 @@
 )
 from deepmd.dpmodel.array_api import (
     Array,
-    support_array_api,
     xp_take_along_axis,
 )
 from deepmd.dpmodel.utils.safe_gradient import (
     safe_for_vector_norm,
 )
 
 
-@support_array_api(version="2023.12")
 def compute_smooth_weight(
     distance: Array,
     rmin: float,
@@ -35,7 +33,6 @@ def compute_smooth_weight(
     return vv
 
 
-@support_array_api(version="2023.12")
 def compute_exp_sw(
     distance: Array,
     rmin: float,

deepmd/dpmodel/utils/network.py

@@ -23,7 +23,6 @@
 )
 from deepmd.dpmodel.array_api import (
     Array,
-    support_array_api,
     xp_add_at,
     xp_bincount,
 )
@@ -259,7 +258,6 @@ def dim_in(self) -> int:
     def dim_out(self) -> int:
         return self.w.shape[1]
 
-    @support_array_api(version="2022.12")
     def call(self, x):  # noqa: ANN001, ANN201
         """Forward pass.
 
@@ -296,7 +294,6 @@ def call(self, x):  # noqa: ANN001, ANN201
         return y
 
 
-@support_array_api(version="2022.12")
 def get_activation_fn(activation_function: str) -> Callable[[np.ndarray], np.ndarray]:
     activation_function = activation_function.lower()
     if activation_function == "tanh":

deepmd/dpmodel/utils/type_embed.py

@@ -8,7 +8,6 @@
 
 from deepmd.dpmodel.array_api import (
     Array,
-    support_array_api,
 )
 from deepmd.dpmodel.common import (
     PRECISION_DICT,
@@ -96,7 +95,6 @@ def __init__(
             trainable=trainable,
         )
 
-    @support_array_api(version="2022.12")
     def call(self) -> Array:
         """Compute the type embedding network."""
         sample_array = self.embedding_net[0]["w"]

source/tests/common/dpmodel/array_api/utils.py

@@ -1,17 +1,10 @@
 # SPDX-License-Identifier: LGPL-3.0-or-later
 import array_api_compat
-from array_api_strict import (
-    set_array_api_strict_flags,
-)
 
 
 class ArrayAPITest:
     """Utils for array API tests."""
 
-    def set_array_api_version(self, func) -> None:
-        """Set the array API version for a function."""
-        set_array_api_strict_flags(api_version=func.array_api_version)
-
     def assert_namespace_equal(self, a, b) -> None:
         """Assert two array has the same namespace."""
         self.assertEqual(