BUG: clip(out=...) is broken (#261)

reviewed at #261

Author: crusaderky

array_api_compat/common/_aliases.py

@@ -12,7 +12,7 @@
 from typing import NamedTuple
 import inspect
 
-from ._helpers import array_namespace, _check_device, device, is_torch_array, is_cupy_namespace
+from ._helpers import array_namespace, _check_device, device, is_cupy_namespace
 
 # These functions are modified from the NumPy versions.
 
@@ -368,23 +368,23 @@ def _isscalar(a):
     if type(max) is int and max >= wrapped_xp.iinfo(x.dtype).max:
         max = None
 
+    dev = device(x)
     if out is None:
-        out = wrapped_xp.asarray(xp.broadcast_to(x, result_shape),
-                                 copy=True, device=device(x))
+        out = wrapped_xp.empty(result_shape, dtype=x.dtype, device=dev)
+    out[()] = x
+
     if min is not None:
-        if is_torch_array(x) and x.dtype == xp.float64 and _isscalar(min):
-            # Avoid loss of precision due to torch defaulting to float32
-            min = wrapped_xp.asarray(min, dtype=xp.float64)
-        a = xp.broadcast_to(wrapped_xp.asarray(min, device=device(x)), result_shape)
+        a = wrapped_xp.asarray(min, dtype=x.dtype, device=dev)
+        a = xp.broadcast_to(a, result_shape)
         ia = (out < a) | xp.isnan(a)
-        # torch requires an explicit cast here
-        out[ia] = wrapped_xp.astype(a[ia], out.dtype)
+        out[ia] = a[ia]
+
     if max is not None:
-        if is_torch_array(x) and x.dtype == xp.float64 and _isscalar(max):
-            max = wrapped_xp.asarray(max, dtype=xp.float64)
-        b = xp.broadcast_to(wrapped_xp.asarray(max, device=device(x)), result_shape)
+        b = wrapped_xp.asarray(max, dtype=x.dtype, device=dev)
+        b = xp.broadcast_to(b, result_shape)
         ib = (out > b) | xp.isnan(b)
-        out[ib] = wrapped_xp.astype(b[ib], out.dtype)
+        out[ib] = b[ib]
+
     # Return a scalar for 0-D
     return out[()]
 

tests/test_common.py

@@ -367,3 +367,18 @@ def test_asarray_copy(library):
         assert all(b[0] == 1.0)
     else:
         assert all(b[0] == 0.0)
+
+
+@pytest.mark.parametrize("library", ["numpy", "cupy", "torch"])
+def test_clip_out(library):
+    """Test non-standard out= parameter for clip()
+
+    (see "Avoid Restricting Behavior that is Outside the Scope of the Standard"
+    in https://data-apis.org/array-api-compat/dev/special-considerations.html)
+    """
+    xp = import_(library, wrapper=True)
+    x = xp.asarray([10, 20, 30])
+    out = xp.zeros_like(x)
+    xp.clip(x, 15, 25, out=out)
+    expect = xp.asarray([15, 20, 25])
+    assert xp.all(out == expect)