Improve clipping time with np_compat.clip

src/array_api_compat/numpy/_aliases.py

@@ -49,7 +49,6 @@
 var = get_xp(np)(_aliases.var)
 cumulative_sum = get_xp(np)(_aliases.cumulative_sum)
 cumulative_prod = get_xp(np)(_aliases.cumulative_prod)
-clip = get_xp(np)(_aliases.clip)
 permute_dims = get_xp(np)(_aliases.permute_dims)
 reshape = get_xp(np)(_aliases.reshape)
 argsort = get_xp(np)(_aliases.argsort)
@@ -106,6 +105,88 @@ def astype(
     return x.astype(dtype=dtype, copy=copy)
 
 
+def clip(
+    x: Array,
+    /,
+    min: float | Array | None = None,
+    max: float | Array | None = None,
+    **kwargs,
+) -> Array:
+    """Array API compatible clip implementation for NumPy.
+
+    NumPy's native ``clip`` is used directly after casting bounds to the
+    input dtype. This keeps the result dtype aligned with ``x.dtype`` and
+    avoids NumPy's default promotion behavior.
+
+    Args:
+        x: Input array.
+        min: Minimum bound. If None, no lower bound is applied.
+        max: Maximum bound. If None, no upper bound is applied.
+        out: Optional output array to store the result, has to have dtype of x
+    """
+    # out is a possible *kwarg for numpy.clip, but not in the array API spec. We handle it here to
+    # avoid having to add it to the array API spec, which would be a breaking change
+    # check if out in kwargs, if so pop it and use it as the out parameter
+    if "out" in kwargs:
+        out = kwargs.pop("out")
+    else:
+        out = None
+
+    def _bound_shape(a: object) -> tuple[int, ...]:
+        if a is None or np.isscalar(a):
+            return ()
+        return np.asarray(a).shape
+
+    dtype = x.dtype
+    out_dtype = out.dtype if out is not None else dtype
+    if out_dtype != dtype:
+        raise ValueError(
+            f"Output array has dtype {out_dtype}, but input array has dtype {dtype}"
+            )
+    min_shape = _bound_shape(min)
+    max_shape = _bound_shape(max)
+
+    # avoid shape broadcasting and copying when not necessary
+    if min_shape == () and max_shape == ():
+        result_shape = x.shape
+    else:
+        result_shape = np.broadcast_shapes(x.shape, min_shape, max_shape)
+
+    # At least handle the case of Python integers correctly.
+    if np.issubdtype(dtype, np.integer):
+        if type(min) is int and min <= np.iinfo(dtype).min:
+            min = None
+        if type(max) is int and max >= np.iinfo(dtype).max:
+            max = None
+
+    if min is None and max is None:
+        if out is None:
+            return x.copy()[()]
+        np.copyto(out, x)
+        return out[()]
+
+    # Cast clip parameters to the input dtype and broadcast them to the result shape.
+    a_min = None
+    if min is not None:
+        a_min = np.asarray(min, dtype=dtype)
+        if a_min.shape != result_shape:
+            # Casting first keeps NumPy from promoting the output dtype.
+            a_min = np.broadcast_to(a_min, result_shape)
+
+    a_max = None
+    if max is not None:
+        a_max = np.asarray(max, dtype=dtype)
+        if a_max.shape != result_shape:
+            # Casting first keeps NumPy from promoting the output dtype.
+            a_max = np.broadcast_to(a_max, result_shape)
+
+    if out is None:
+        out = np.empty(result_shape, dtype=dtype)
+
+    np.clip(x, a_min, a_max, out=out, casting="no", **kwargs)
+    return out[()]
+
+
 # count_nonzero returns a python int for axis=None and keepdims=False
 # https://github.com/numpy/numpy/issues/17562
 def count_nonzero(
@@ -115,7 +196,9 @@ def count_nonzero(
 ) -> Array:
     # NOTE: this is currently incorrectly typed in numpy, but will be fixed in
     # numpy 2.2.5 and 2.3.0: https://github.com/numpy/numpy/pull/28750
-    result = cast("Any", np.count_nonzero(x, axis=axis, keepdims=keepdims))  # pyright: ignore[reportArgumentType, reportCallIssue]
+    result = cast(
+        "Any", np.count_nonzero(x, axis=axis, keepdims=keepdims)
+    )  # pyright: ignore[reportArgumentType, reportCallIssue]
     if axis is None and not keepdims:
         return np.asarray(result)
     return result
@@ -128,20 +211,21 @@ def take_along_axis(x: Array, indices: Array, /, *, axis: int = -1) -> Array:
 
 # ceil, floor, and trunc return integers for integer inputs in NumPy < 2
 
+
 def ceil(x: Array, /) -> Array:
-    if np.__version__ < '2' and np.issubdtype(x.dtype, np.integer):
+    if np.__version__ < "2" and np.issubdtype(x.dtype, np.integer):
         return x.copy()
     return np.ceil(x)
 
 
 def floor(x: Array, /) -> Array:
-    if np.__version__ < '2' and np.issubdtype(x.dtype, np.integer):
+    if np.__version__ < "2" and np.issubdtype(x.dtype, np.integer):
         return x.copy()
     return np.floor(x)
 
 
 def trunc(x: Array, /) -> Array:
-    if np.__version__ < '2' and np.issubdtype(x.dtype, np.integer):
+    if np.__version__ < "2" and np.issubdtype(x.dtype, np.integer):
         return x.copy()
     return np.trunc(x)
 
@@ -173,6 +257,7 @@ def trunc(x: Array, /) -> Array:
     "atan",
     "atan2",
     "atanh",
+    "clip",
     "ceil",
     "floor",
     "trunc",
@@ -183,7 +268,7 @@ def trunc(x: Array, /) -> Array:
     "concat",
     "count_nonzero",
     "pow",
-    "take_along_axis"
+    "take_along_axis",
 ]
 
 

tests/test_numpy.py

@@ -1,5 +1,5 @@
-"""Test "unspecified" behavior which we cannot easily test in the Array API test suite.
-"""
+"""Test "unspecified" behavior which we cannot easily test in the Array API test suite."""
+
 import warnings
 import pytest
 
@@ -10,6 +10,84 @@
 
 from array_api_compat import is_array_api_obj
 
+
+def test_numpy_clip_out_and_broadcast():
+    from array_api_compat import numpy as xp
+
+    x = xp.asarray([[10, 20, 30], [40, 50, 60]], dtype=np.uint8)
+    min_bound = xp.asarray([15, 35, 55], dtype=np.int16)
+    max_bound = xp.asarray([25, 45, 65], dtype=np.int16)
+    out = xp.empty_like(x)
+
+    result = xp.clip(x, min_bound, max_bound, out=out)
+
+    np.testing.assert_array_equal(result, xp.asarray([[15, 35, 55], [25, 45, 60]], dtype=np.uint8))
+    assert result.dtype == x.dtype
+    np.testing.assert_array_equal(out, xp.asarray([[15, 35, 55], [25, 45, 60]], dtype=np.uint8))
+
+
+def test_numpy_clip_uint8_casts_bounds_outside_range():
+    from array_api_compat import numpy as xp
+
+    x = xp.asarray([0, 10, 250], dtype=np.uint8)
+    min_bound = np.int16(-1)
+    max_bound = np.int16(200)
+
+    result = xp.clip(x, min_bound, max_bound)
+
+    assert result.dtype == x.dtype
+    np.testing.assert_array_equal(result, xp.asarray([200, 200, 200], dtype=np.uint8))
+
+
+def test_numpy_clip_int64_casts_bounds_outside_range():
+    from array_api_compat import numpy as xp
+
+    x = xp.asarray([-(2**63), -1, 0, 2**63 - 1], dtype=np.int64)
+    min_bound = np.float64(-1e20)
+    max_bound = np.float64(1e20)
+
+    result = xp.clip(x, min_bound, max_bound)
+
+    assert result.dtype == x.dtype
+    np.testing.assert_array_equal(
+        result,
+        xp.asarray(
+            [
+                np.iinfo(np.int64).min,
+                np.iinfo(np.int64).min,
+                np.iinfo(np.int64).min,
+                np.iinfo(np.int64).min,
+            ],
+            dtype=np.int64,
+        ),
+    )
+
+
+def test_numpy_clip_float16_casts_bounds_outside_range():
+    from array_api_compat import numpy as xp
+
+    x = xp.asarray([0.0, 1.5, 3.0], dtype=np.float16)
+    min_bound = np.float32(-1e10)
+    max_bound = np.float32(2.0)
+
+    result = xp.clip(x, min_bound, max_bound)
+
+    assert result.dtype == x.dtype
+    np.testing.assert_array_equal(result, xp.asarray([0.0, 1.5, 2.0], dtype=np.float16))
+
+
+def test_numpy_clip_returns_copy_when_unbounded():
+    from array_api_compat import numpy as xp
+
+    x = xp.arange(8, dtype=np.int64)
+
+    y = xp.clip(x)
+
+    assert y.dtype == x.dtype
+    assert not np.shares_memory(x, y)
+    np.testing.assert_array_equal(y, x)
+
+
 def test_matrix_is_not_array_api_obj():
     assert is_array_api_obj(np.asarray(3))
     assert is_array_api_obj(np.float64(3))