Merge branch 'main' into more-2023

Author: asmeurer

.github/workflows/publish-package.yml

@@ -94,7 +94,7 @@ jobs:
       if: >-
         (github.event_name == 'push' && startsWith(github.ref, 'refs/tags'))
         || (github.event_name == 'workflow_dispatch' && github.event.inputs.publish == 'true')
-      uses: pypa/gh-action-pypi-publish@v1.9.0
+      uses: pypa/gh-action-pypi-publish@v1.10.1
       with:
         repository-url: https://test.pypi.org/legacy/
         print-hash: true
@@ -107,6 +107,6 @@ jobs:
 
     - name: Publish distribution 📦 to PyPI
       if: github.event_name == 'push' && startsWith(github.ref, 'refs/tags')
-      uses: pypa/gh-action-pypi-publish@v1.9.0
+      uses: pypa/gh-action-pypi-publish@v1.10.1
       with:
         print-hash: true

array_api_compat/common/_helpers.py

@@ -202,7 +202,6 @@ def is_jax_array(x):
 
     return isinstance(x, jax.Array) or _is_jax_zero_gradient_array(x)
 
-
 def is_pydata_sparse_array(x) -> bool:
     """
     Return True if `x` is an array from the `sparse` package.
@@ -255,6 +254,166 @@ def is_array_api_obj(x):
         or is_pydata_sparse_array(x) \
         or hasattr(x, '__array_namespace__')
 
+def _compat_module_name():
+    assert __name__.endswith('.common._helpers')
+    return __name__.removesuffix('.common._helpers')
+
+def is_numpy_namespace(xp) -> bool:
+    """
+    Returns True if `xp` is a NumPy namespace.
+
+    This includes both NumPy itself and the version wrapped by array-api-compat.
+
+    See Also
+    --------
+
+    array_namespace
+    is_cupy_namespace
+    is_torch_namespace
+    is_ndonnx_namespace
+    is_dask_namespace
+    is_jax_namespace
+    is_pydata_sparse_namespace
+    is_array_api_strict_namespace
+    """
+    return xp.__name__ in {'numpy', _compat_module_name() + '.numpy'}
+
+def is_cupy_namespace(xp) -> bool:
+    """
+    Returns True if `xp` is a CuPy namespace.
+
+    This includes both CuPy itself and the version wrapped by array-api-compat.
+
+    See Also
+    --------
+
+    array_namespace
+    is_numpy_namespace
+    is_torch_namespace
+    is_ndonnx_namespace
+    is_dask_namespace
+    is_jax_namespace
+    is_pydata_sparse_namespace
+    is_array_api_strict_namespace
+    """
+    return xp.__name__ in {'cupy', _compat_module_name() + '.cupy'}
+
+def is_torch_namespace(xp) -> bool:
+    """
+    Returns True if `xp` is a PyTorch namespace.
+
+    This includes both PyTorch itself and the version wrapped by array-api-compat.
+
+    See Also
+    --------
+
+    array_namespace
+    is_numpy_namespace
+    is_cupy_namespace
+    is_ndonnx_namespace
+    is_dask_namespace
+    is_jax_namespace
+    is_pydata_sparse_namespace
+    is_array_api_strict_namespace
+    """
+    return xp.__name__ in {'torch', _compat_module_name() + '.torch'}
+    
+
+def is_ndonnx_namespace(xp):
+    """
+    Returns True if `xp` is an NDONNX namespace.
+
+    See Also
+    --------
+
+    array_namespace
+    is_numpy_namespace
+    is_cupy_namespace
+    is_torch_namespace
+    is_dask_namespace
+    is_jax_namespace
+    is_pydata_sparse_namespace
+    is_array_api_strict_namespace
+    """
+    return xp.__name__ == 'ndonnx'
+
+def is_dask_namespace(xp):
+    """
+    Returns True if `xp` is a Dask namespace.
+
+    This includes both ``dask.array`` itself and the version wrapped by array-api-compat.
+
+    See Also
+    --------
+
+    array_namespace
+    is_numpy_namespace
+    is_cupy_namespace
+    is_torch_namespace
+    is_ndonnx_namespace
+    is_jax_namespace
+    is_pydata_sparse_namespace
+    is_array_api_strict_namespace
+    """
+    return xp.__name__ in {'dask.array', _compat_module_name() + '.dask.array'}
+
+def is_jax_namespace(xp):
+    """
+    Returns True if `xp` is a JAX namespace.
+
+    This includes ``jax.numpy`` and ``jax.experimental.array_api`` which existed in
+    older versions of JAX.
+
+    See Also
+    --------
+
+    array_namespace
+    is_numpy_namespace
+    is_cupy_namespace
+    is_torch_namespace
+    is_ndonnx_namespace
+    is_dask_namespace
+    is_pydata_sparse_namespace
+    is_array_api_strict_namespace
+    """
+    return xp.__name__ in {'jax.numpy', 'jax.experimental.array_api'}
+
+def is_pydata_sparse_namespace(xp):
+    """
+    Returns True if `xp` is a pydata/sparse namespace.
+
+    See Also
+    --------
+
+    array_namespace
+    is_numpy_namespace
+    is_cupy_namespace
+    is_torch_namespace
+    is_ndonnx_namespace
+    is_dask_namespace
+    is_jax_namespace
+    is_array_api_strict_namespace
+    """
+    return xp.__name__ == 'sparse'
+
+def is_array_api_strict_namespace(xp):
+    """
+    Returns True if `xp` is an array-api-strict namespace.
+
+    See Also
+    --------
+
+    array_namespace
+    is_numpy_namespace
+    is_cupy_namespace
+    is_torch_namespace
+    is_ndonnx_namespace
+    is_dask_namespace
+    is_jax_namespace
+    is_pydata_sparse_namespace
+    """
+    return xp.__name__ == 'array_api_strict'
+
 def _check_api_version(api_version):
     if api_version == '2021.12':
         warnings.warn("The 2021.12 version of the array API specification was requested but the returned namespace is actually version 2022.12")
@@ -643,13 +802,21 @@ def size(x):
     "device",
     "get_namespace",
     "is_array_api_obj",
+    "is_array_api_strict_namespace",
     "is_cupy_array",
+    "is_cupy_namespace",
     "is_dask_array",
+    "is_dask_namespace",
     "is_jax_array",
+    "is_jax_namespace",
     "is_numpy_array",
+    "is_numpy_namespace",
     "is_torch_array",
+    "is_torch_namespace",
     "is_ndonnx_array",
+    "is_ndonnx_namespace",
     "is_pydata_sparse_array",
+    "is_pydata_sparse_namespace",
     "size",
     "to_device",
 ]

array_api_compat/cupy/_aliases.py

@@ -108,6 +108,13 @@ def asarray(
 
         return cp.array(obj, dtype=dtype, **kwargs)
 
+def sign(x: ndarray, /) -> ndarray:
+    # CuPy sign() does not propagate nans. See
+    # https://github.com/data-apis/array-api-compat/issues/136
+    out = cp.sign(x)
+    out[cp.isnan(x)] = cp.nan
+    return out
+
 # These functions are completely new here. If the library already has them
 # (i.e., numpy 2.0), use the library version instead of our wrapper.
 if hasattr(cp, 'vecdot'):
@@ -129,6 +136,6 @@ def asarray(
                               'acos', 'acosh', 'asin', 'asinh', 'atan',
                               'atan2', 'atanh', 'bitwise_left_shift',
                               'bitwise_invert', 'bitwise_right_shift',
-                              'concat', 'pow']
+                              'concat', 'pow', 'sign']
 
 _all_ignore = ['cp', 'get_xp']

array_api_compat/torch/_aliases.py

@@ -709,6 +709,21 @@ def take(x: array, indices: array, /, *, axis: Optional[int] = None, **kwargs) -
         axis = 0
     return torch.index_select(x, axis, indices, **kwargs)
 
+def sign(x: array, /) -> array:
+    # torch sign() does not support complex numbers and does not propagate
+    # nans. See https://github.com/data-apis/array-api-compat/issues/136
+    if x.dtype.is_complex:
+        out = x/torch.abs(x)
+        # sign(0) = 0 but the above formula would give nan
+        out[x == 0+0j] = 0+0j
+        return out
+    else:
+        out = torch.sign(x)
+        if x.dtype.is_floating_point:
+            out[torch.isnan(x)] = torch.nan
+        return out
+
+
 __all__ = ['result_type', 'can_cast', 'permute_dims', 'bitwise_invert',
            'newaxis', 'conj', 'add', 'atan2', 'bitwise_and', 'bitwise_left_shift',
            'bitwise_or', 'bitwise_right_shift', 'bitwise_xor', 'copysign',
@@ -722,6 +737,6 @@ def take(x: array, indices: array, /, *, axis: Optional[int] = None, **kwargs) -
            'broadcast_arrays', 'UniqueAllResult', 'UniqueCountsResult',
            'UniqueInverseResult', 'unique_all', 'unique_counts',
            'unique_inverse', 'unique_values', 'matmul', 'matrix_transpose',
-           'vecdot', 'tensordot', 'isdtype', 'take']
+           'vecdot', 'tensordot', 'isdtype', 'take', 'sign']
 
 _all_ignore = ['torch', 'get_xp']

cupy-xfails.txt

@@ -160,7 +160,6 @@ array_api_tests/test_special_cases.py::test_unary[expm1(x_i is -0) -> -0]
 array_api_tests/test_special_cases.py::test_unary[floor(x_i is -0) -> -0]
 array_api_tests/test_special_cases.py::test_unary[log1p(x_i is -0) -> -0]
 array_api_tests/test_special_cases.py::test_unary[round(x_i is -0) -> -0]
-array_api_tests/test_special_cases.py::test_unary[sign(x_i is NaN) -> NaN]
 array_api_tests/test_special_cases.py::test_unary[sin(x_i is -0) -> -0]
 array_api_tests/test_special_cases.py::test_unary[sinh(x_i is -0) -> -0]
 array_api_tests/test_special_cases.py::test_unary[sqrt(x_i is -0) -> -0]

docs/helper-functions.rst

@@ -40,7 +40,7 @@ instead, which would be wrapped.
 Inspection Helpers
 ------------------
 
-These convenience functions can be used to test if an array comes from a
+These convenience functions can be used to test if an array or namespace comes from a
 specific library without importing that library if it hasn't been imported
 yet.
 
@@ -51,3 +51,11 @@ yet.
 .. autofunction:: is_jax_array
 .. autofunction:: is_pydata_sparse_array
 .. autofunction:: is_ndonnx_array
+.. autofunction:: is_numpy_namespace
+.. autofunction:: is_cupy_namespace
+.. autofunction:: is_torch_namespace
+.. autofunction:: is_dask_namespace
+.. autofunction:: is_jax_namespace
+.. autofunction:: is_pydata_sparse_namespace
+.. autofunction:: is_ndonnx_namespace
+.. autofunction:: is_array_api_strict_namespace

tests/test_common.py

@@ -1,5 +1,9 @@
-from array_api_compat import (is_numpy_array, is_cupy_array, is_torch_array, # noqa: F401
-                              is_dask_array, is_jax_array, is_pydata_sparse_array)
+from array_api_compat import (  # noqa: F401
+    is_numpy_array, is_cupy_array, is_torch_array,
+    is_dask_array, is_jax_array, is_pydata_sparse_array,
+    is_numpy_namespace, is_cupy_namespace, is_torch_namespace,
+    is_dask_namespace, is_jax_namespace, is_pydata_sparse_namespace,
+)
 
 from array_api_compat import is_array_api_obj, device, to_device
 
@@ -10,7 +14,7 @@
 import array
 from numpy.testing import assert_allclose
 
-is_functions = {
+is_array_functions = {
     'numpy': 'is_numpy_array',
     'cupy': 'is_cupy_array',
     'torch': 'is_torch_array',
@@ -19,18 +23,38 @@
     'sparse': 'is_pydata_sparse_array',
 }
 
-@pytest.mark.parametrize('library', is_functions.keys())
-@pytest.mark.parametrize('func', is_functions.values())
+is_namespace_functions = {
+    'numpy': 'is_numpy_namespace',
+    'cupy': 'is_cupy_namespace',
+    'torch': 'is_torch_namespace',
+    'dask.array': 'is_dask_namespace',
+    'jax.numpy': 'is_jax_namespace',
+    'sparse': 'is_pydata_sparse_namespace',
+}
+
+
+@pytest.mark.parametrize('library', is_array_functions.keys())
+@pytest.mark.parametrize('func', is_array_functions.values())
 def test_is_xp_array(library, func):
     lib = import_(library)
     is_func = globals()[func]
 
     x = lib.asarray([1, 2, 3])
 
-    assert is_func(x) == (func == is_functions[library])
+    assert is_func(x) == (func == is_array_functions[library])
 
     assert is_array_api_obj(x)
 
+
+@pytest.mark.parametrize('library', is_namespace_functions.keys())
+@pytest.mark.parametrize('func', is_namespace_functions.values())
+def test_is_xp_namespace(library, func):
+    lib = import_(library)
+    is_func = globals()[func]
+
+    assert is_func(lib) == (func == is_namespace_functions[library])
+
+
 @pytest.mark.parametrize("library", all_libraries)
 def test_device(library):
     xp = import_(library, wrapper=True)
@@ -64,8 +88,8 @@ def test_to_device_host(library):
     assert_allclose(x, expected)
 
 
-@pytest.mark.parametrize("target_library", is_functions.keys())
-@pytest.mark.parametrize("source_library", is_functions.keys())
+@pytest.mark.parametrize("target_library", is_array_functions.keys())
+@pytest.mark.parametrize("source_library", is_array_functions.keys())
 def test_asarray_cross_library(source_library, target_library, request):
     if source_library == "dask.array" and target_library == "torch":
         # Allow rest of test to execute instead of immediately xfailing
@@ -81,7 +105,7 @@ def test_asarray_cross_library(source_library, target_library, request):
         pytest.skip(reason="`sparse` does not allow implicit densification")
     src_lib = import_(source_library, wrapper=True)
     tgt_lib = import_(target_library, wrapper=True)
-    is_tgt_type = globals()[is_functions[target_library]]
+    is_tgt_type = globals()[is_array_functions[target_library]]
 
     a = src_lib.asarray([1, 2, 3])
     b = tgt_lib.asarray(a)
@@ -96,7 +120,7 @@ def test_asarray_copy(library):
     # should be able to delete this.
     xp = import_(library, wrapper=True)
     asarray = xp.asarray
-    is_lib_func = globals()[is_functions[library]]
+    is_lib_func = globals()[is_array_functions[library]]
     all = xp.all if library != 'dask.array' else lambda x: xp.all(x).compute()
 
     if library == 'numpy' and xp.__version__[0] < '2' and not hasattr(xp, '_CopyMode') :

tests/test_vendoring.py

@@ -20,6 +20,7 @@ def test_vendoring_torch():
 
     uses_torch._test_torch()
 
+
 def test_vendoring_dask():
     from vendor_test import uses_dask
     uses_dask._test_dask()