Wrap clip() for torch

Author: asmeurer

array_api_compat/common/_aliases.py

@@ -315,11 +315,12 @@ def _isscalar(a):
     if min is not None:
         a = xp.broadcast_to(xp.asarray(min), result_shape)
         ia = (out < a) | xp.isnan(a)
-        out[ia] = a[ia]
+        # torch requires an explicit cast here
+        out[ia] = wrapped_xp.astype(a[ia], out.dtype)
     if max is not None:
         b = xp.broadcast_to(xp.asarray(max), result_shape)
         ib = (out > b) | xp.isnan(b)
-        out[ib] = b[ib]
+        out[ib] = wrapped_xp.astype(b[ib], out.dtype)
     # Return a scalar for 0-D
     return out[()]
 

array_api_compat/torch/_aliases.py

@@ -4,7 +4,7 @@
 from builtins import all as _builtin_all, any as _builtin_any
 
 from ..common._aliases import (matrix_transpose as _aliases_matrix_transpose,
-                               vecdot as _aliases_vecdot)
+                               vecdot as _aliases_vecdot, clip as _aliases_clip)
 from .._internal import get_xp
 
 import torch
@@ -189,6 +189,8 @@ def min(x: array, /, *, axis: Optional[Union[int, Tuple[int, ...]]] = None, keep
         return torch.clone(x)
     return torch.amin(x, axis, keepdims=keepdims)
 
+clip = get_xp(torch)(_aliases_clip)
+
 # torch.sort also returns a tuple
 # https://github.com/pytorch/pytorch/issues/70921
 def sort(x: array, /, *, axis: int = -1, descending: bool = False, stable: bool = True, **kwargs) -> array:
@@ -706,8 +708,8 @@ def take(x: array, indices: array, /, *, axis: Optional[int] = None, **kwargs) -
            'bitwise_or', 'bitwise_right_shift', 'bitwise_xor', 'copysign',
            'divide', 'equal', 'floor_divide', 'greater', 'greater_equal',
            'less', 'less_equal', 'logaddexp', 'multiply', 'not_equal', 'pow',
-           'remainder', 'subtract', 'max', 'min', 'sort', 'prod', 'sum',
-           'any', 'all', 'mean', 'std', 'var', 'concat', 'squeeze',
+           'remainder', 'subtract', 'max', 'min', 'clip', 'sort', 'prod',
+           'sum', 'any', 'all', 'mean', 'std', 'var', 'concat', 'squeeze',
            'broadcast_to', 'flip', 'roll', 'nonzero', 'where', 'reshape',
            'arange', 'eye', 'linspace', 'full', 'ones', 'zeros', 'empty',
            'tril', 'triu', 'expand_dims', 'astype', 'broadcast_arrays',