diff --git a/src/array_api_compat/numpy/_aliases.py b/src/array_api_compat/numpy/_aliases.py index 87b3c2f3..6ef3060c 100644 --- a/src/array_api_compat/numpy/_aliases.py +++ b/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", ] diff --git a/tests/test_numpy.py b/tests/test_numpy.py index a139d428..de8243f1 100644 --- a/tests/test_numpy.py +++ b/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))