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,112 @@ 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.
+        **kwargs: Additional keyword arguments passed to ``np.clip``.
+        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)
+
+    # Handle cases where the bounds are outside the range of the integer input dtype
+    # this covers integer arrays for float and integer bounds
+    # Also handle cases where the min/max are arrays/lists (replace values below min with iinfo.min and above max with iinfo.max)
+    if np.issubdtype(dtype, np.integer):
+
+        if np.issubdtype(type(min), np.integer) and min <= np.iinfo(dtype).min:
+            min = None
+        elif np.issubdtype(type(min), np.floating) and min < np.iinfo(dtype).min:
+            min = np.iinfo(dtype).min
+        elif isinstance(min, Array):
+            min[min < np.iinfo(dtype).min] = np.iinfo(dtype).min
+        elif isinstance(min, (list,tuple)):
+            min = np.asarray(min)
+            min[min < np.iinfo(dtype).min] = np.iinfo(dtype).min
+
+
+        if np.issubdtype(type(max), np.integer) and max >= np.iinfo(dtype).max:
+            max = None
+        elif np.issubdtype(type(max), np.floating) and max > np.iinfo(dtype).max:
+            max = np.iinfo(dtype).max
+        elif isinstance(max, Array):
+            max[max > np.iinfo(dtype).max] = np.iinfo(dtype).max
+        elif isinstance(max, (list,tuple)):
+            max = np.asarray(max)
+            max[max > np.iinfo(dtype).max] = np.iinfo(dtype).max
+
+    # In the case of downcasting floats numpy replaces out of bounds with inf 
+    # This automatically handles those cases
+
+    # Early return for simple cases
+    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(
@@ -173,6 +278,7 @@ def trunc(x: Array, /) -> Array:
     "atan",
     "atan2",
     "atanh",
+    "clip",
     "ceil",
     "floor",
     "trunc",
@@ -183,7 +289,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
 
@@ -9,6 +9,128 @@
     pytestmark = pytest.skip(allow_module_level=True, reason="numpy not found")
 
 from array_api_compat import is_array_api_obj
+from array_api_compat import numpy as xp
+
+def test_numpy_clip_out_and_broadcast():
+
+    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_all_bounds_work_with_int_arrays():
+    """Test that integer bounds outside the range of input dtype still work for integer arrays"""
+    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([0, 10, 200], dtype=np.uint8))
+
+
+    # min and max bounds are below what can be represented by int64, 
+    # so they should be clipped to the min/max of int64
+    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(
+            [
+            -(2**63), -1, 0, 2**63 - 1
+            ],
+            dtype=np.int64,
+        ),
+    )
+
+
+def test_array_min_max_broadcasting_when_clipped():
+    """ Tests a min as tuple list array of floats input for an integer array
+    Should be clipped by min/max of the integer array"""
+    x=xp.asarray([0, 10, 100], dtype=np.uint8)
+    min_bound = xp.asarray([np.float64(-1e20), 5.0, 200.0], dtype=np.float32)
+    max_bound = None
+    result=xp.clip(x, min_bound, max_bound)
+    assert result.dtype == x.dtype
+    np.testing.assert_array_equal(result, xp.asarray([0, 10, 200], dtype=np.uint8))
+
+    # now test with a tuple
+    min_bound = (np.float64(-1e20), 5.0, 200.0)
+    result=xp.clip(x, min_bound, max_bound)
+    assert result.dtype == x.dtype
+    np.testing.assert_array_equal(result, xp.asarray([0, 10, 200], dtype=np.uint8))
+
+    # test with a list
+    min_bound = [np.float64(-1e20), 5.0, 200.0]
+    result=xp.clip(x, min_bound, max_bound)
+    assert result.dtype == x.dtype
+    np.testing.assert_array_equal(result, xp.asarray([0, 10, 200], dtype=np.uint8))
+
+def test_numpy_type_promotion():
+    """ Added to address comment from main alias file:
+    # np.clip does type promotion but the array API clip requires that the
+    # output have the same dtype as x. We do this instead of just downcasting
+    # the result of xp.clip() to handle some corner cases better (e.g.,
+    # avoiding uint64 -> float64 promotion).
+    """
+    # ensure clipping with float bounds
+    x = xp.asarray([-(2**63), -1, 0, 2**63 - 1], dtype=np.int64)
+    min_bound = np.float64(-1.0001)
+    max_bound = np.float64(1.0001)
+    result = xp.clip(x, min_bound, max_bound)
+    assert result.dtype == x.dtype
+    np.testing.assert_array_equal(result, xp.asarray([-1, -1, 0, 1], dtype=np.int64))
+
+    # ensure clipping with int16 bounds
+    min_bound = np.int16(-1)
+    max_bound = np.int16(1)
+    result = xp.clip(x, min_bound, max_bound)
+    assert result.dtype == x.dtype
+    np.testing.assert_array_equal(result, xp.asarray([-1, -1, 0, 1], dtype=np.int64))
+
+    # final test with uint8 image and int64 bounds
+    x = xp.asarray([0, 10, 250], dtype=np.uint8)
+    min_bound = np.int64(-1)
+    max_bound = np.int64(32)
+    result = xp.clip(x, min_bound, max_bound)
+    assert result.dtype == x.dtype
+    np.testing.assert_array_equal(result, xp.asarray([0, 10, 32], dtype=np.uint8))
+
+def test_numpy_clip_float16_casts_bounds_outside_range():
+    """Test that float16 bounds outside the range of input dtype still work for float16 arrays"""
+    x = xp.asarray([0.0, 1.5, 3.0], dtype=np.float16)
+    min_bound = np.float32(-1e10) # outside of float16 range
+    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():
+
+    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))