diff --git a/src/array_api_compat/numpy/_aliases.py b/src/array_api_compat/numpy/_aliases.py index 87b3c2f3..fe7912f6 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,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", ] diff --git a/tests/test_numpy.py b/tests/test_numpy.py index a139d428..9af4e3b4 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 @@ -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))