Add a wrapper for sign for NumPy-likes

Fixes #183
data-apis · Oct 24, 2024 · dd44814 · dd44814
1 parent 5affae5
commit dd44814
Show file tree

Hide file tree

Showing 4 changed files with 15 additions and 2 deletions.
diff --git a/array_api_compat/common/_aliases.py b/array_api_compat/common/_aliases.py
@@ -530,11 +530,22 @@ def unstack(x: ndarray, /, xp, *, axis: int = 0) -> Tuple[ndarray, ...]:
         raise ValueError("Input array must be at least 1-d.")
     return tuple(xp.moveaxis(x, axis, 0))
 
+# numpy 1.26 does not use the standard definition for sign on complex numbers
+
+def sign(x: array, /, xp, **kwargs) -> array:
+    if isdtype(x.dtype, 'complex floating', xp=xp):
+        out = (x/xp.abs(x, **kwargs))[...]
+        # sign(0) = 0 but the above formula would give nan
+        out[x == 0+0j] = 0+0j
+        return out[()]
+    else:
+        return xp.sign(x, **kwargs)
+
 __all__ = ['arange', 'empty', 'empty_like', 'eye', 'full', 'full_like',
            'linspace', 'ones', 'ones_like', 'zeros', 'zeros_like',
            'UniqueAllResult', 'UniqueCountsResult', 'UniqueInverseResult',
            'unique_all', 'unique_counts', 'unique_inverse', 'unique_values',
            'astype', 'std', 'var', 'cumulative_sum', 'clip', 'permute_dims',
            'reshape', 'argsort', 'sort', 'nonzero', 'ceil', 'floor', 'trunc',
            'matmul', 'matrix_transpose', 'tensordot', 'vecdot', 'isdtype',
-           'unstack']
+           'unstack', 'sign']
diff --git a/array_api_compat/cupy/_aliases.py b/array_api_compat/cupy/_aliases.py
@@ -62,6 +62,7 @@
 matmul = get_xp(cp)(_aliases.matmul)
 matrix_transpose = get_xp(cp)(_aliases.matrix_transpose)
 tensordot = get_xp(cp)(_aliases.tensordot)
+sign = get_xp(cp)(_aliases.sign)
 
 _copy_default = object()
 

diff --git a/array_api_compat/dask/array/_aliases.py b/array_api_compat/dask/array/_aliases.py
@@ -104,7 +104,7 @@ def _dask_arange(
 trunc = get_xp(np)(_aliases.trunc)
 matmul = get_xp(np)(_aliases.matmul)
 tensordot = get_xp(np)(_aliases.tensordot)
-
+sign = get_xp(np)(_aliases.sign)
 
 # asarray also adds the copy keyword, which is not present in numpy 1.0.
 def asarray(

diff --git a/array_api_compat/numpy/_aliases.py b/array_api_compat/numpy/_aliases.py
@@ -62,6 +62,7 @@
 matmul = get_xp(np)(_aliases.matmul)
 matrix_transpose = get_xp(np)(_aliases.matrix_transpose)
 tensordot = get_xp(np)(_aliases.tensordot)
+sign = get_xp(np)(_aliases.sign)
 
 def _supports_buffer_protocol(obj):
     try: