array-api-compat 1.6__tar.gz → 1.7__tar.gz

{array_api_compat-1.6 → array_api_compat-1.7}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.1
 Name: array_api_compat
-Version: 1.6
+Version: 1.7
 Summary: A wrapper around NumPy and other array libraries to make them compatible with the Array API standard
 Home-page: https://data-apis.org/array-api-compat/
 Author: Consortium for Python Data API Standards
@@ -23,13 +23,15 @@ Provides-Extra: pytorch
 Requires-Dist: pytorch; extra == "pytorch"
 Provides-Extra: dask
 Requires-Dist: dask; extra == "dask"
+Provides-Extra: sprase
+Requires-Dist: sparse>=0.15.1; extra == "sprase"
 
 # Array API compatibility library
 
 This is a small wrapper around common array libraries that is compatible with
 the [Array API standard](https://data-apis.org/array-api/latest/). Currently,
-NumPy, CuPy, PyTorch, Dask, and JAX are supported. If you want support for other array
-libraries, or if you encounter any issues, please [open an
+NumPy, CuPy, PyTorch, Dask, JAX and `sparse` are supported. If you want support
+for other array libraries, or if you encounter any issues, please [open an
 issue](https://github.com/data-apis/array-api-compat/issues).
 
 See the documentation for more details https://data-apis.org/array-api-compat/

{array_api_compat-1.6 → array_api_compat-1.7}/README.md

@@ -2,8 +2,8 @@
 
 This is a small wrapper around common array libraries that is compatible with
 the [Array API standard](https://data-apis.org/array-api/latest/). Currently,
-NumPy, CuPy, PyTorch, Dask, and JAX are supported. If you want support for other array
-libraries, or if you encounter any issues, please [open an
+NumPy, CuPy, PyTorch, Dask, JAX and `sparse` are supported. If you want support
+for other array libraries, or if you encounter any issues, please [open an
 issue](https://github.com/data-apis/array-api-compat/issues).
 
 See the documentation for more details https://data-apis.org/array-api-compat/

{array_api_compat-1.6 → array_api_compat-1.7}/array_api_compat/__init__.py

@@ -17,6 +17,6 @@ to ensure they are not using functionality outside of the standard, but prefer
 this implementation for the default when working with NumPy arrays.
 
 """
-__version__ = '1.6'
+__version__ = '1.7'
 
 from .common import *  # noqa: F401, F403

{array_api_compat-1.6 → array_api_compat-1.7}/array_api_compat/common/_helpers.py

@@ -18,6 +18,20 @@ import math
 import inspect
 import warnings
 
+def _is_jax_zero_gradient_array(x):
+    """Return True if `x` is a zero-gradient array.
+
+    These arrays are a design quirk of Jax that may one day be removed.
+    See https://github.com/google/jax/issues/20620.
+    """
+    if 'numpy' not in sys.modules or 'jax' not in sys.modules:
+        return False
+
+    import numpy as np
+    import jax
+
+    return isinstance(x, np.ndarray) and x.dtype == jax.float0
+
 def is_numpy_array(x):
     """
     Return True if `x` is a NumPy array.
@@ -36,6 +50,7 @@ def is_numpy_array(x):
     is_torch_array
     is_dask_array
     is_jax_array
+    is_pydata_sparse_array
     """
     # Avoid importing NumPy if it isn't already
     if 'numpy' not in sys.modules:
@@ -44,7 +59,8 @@ def is_numpy_array(x):
     import numpy as np
 
     # TODO: Should we reject ndarray subclasses?
-    return isinstance(x, (np.ndarray, np.generic))
+    return (isinstance(x, (np.ndarray, np.generic))
+            and not _is_jax_zero_gradient_array(x))
 
 def is_cupy_array(x):
     """
@@ -64,6 +80,7 @@ def is_cupy_array(x):
     is_torch_array
     is_dask_array
     is_jax_array
+    is_pydata_sparse_array
     """
     # Avoid importing NumPy if it isn't already
     if 'cupy' not in sys.modules:
@@ -90,6 +107,7 @@ def is_torch_array(x):
     is_cupy_array
     is_dask_array
     is_jax_array
+    is_pydata_sparse_array
     """
     # Avoid importing torch if it isn't already
     if 'torch' not in sys.modules:
@@ -116,6 +134,7 @@ def is_dask_array(x):
     is_cupy_array
     is_torch_array
     is_jax_array
+    is_pydata_sparse_array
     """
     # Avoid importing dask if it isn't already
     if 'dask.array' not in sys.modules:
@@ -142,6 +161,7 @@ def is_jax_array(x):
     is_cupy_array
     is_torch_array
     is_dask_array
+    is_pydata_sparse_array
     """
     # Avoid importing jax if it isn't already
     if 'jax' not in sys.modules:
@@ -149,7 +169,36 @@ def is_jax_array(x):
 
     import jax
 
-    return isinstance(x, jax.Array)
+    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.
+
+    This function does not import `sparse` if it has not already been imported
+    and is therefore cheap to use.
+
+
+    See Also
+    --------
+
+    array_namespace
+    is_array_api_obj
+    is_numpy_array
+    is_cupy_array
+    is_torch_array
+    is_dask_array
+    is_jax_array
+    """
+    # Avoid importing jax if it isn't already
+    if 'sparse' not in sys.modules:
+        return False
+
+    import sparse
+
+    # TODO: Account for other backends.
+    return isinstance(x, sparse.SparseArray)
 
 def is_array_api_obj(x):
     """
@@ -170,6 +219,7 @@ def is_array_api_obj(x):
         or is_torch_array(x) \
         or is_dask_array(x) \
         or is_jax_array(x) \
+        or is_pydata_sparse_array(x) \
         or hasattr(x, '__array_namespace__')
 
 def _check_api_version(api_version):
@@ -238,6 +288,7 @@ def array_namespace(*xs, api_version=None, use_compat=None):
     is_torch_array
     is_dask_array
     is_jax_array
+    is_pydata_sparse_array
 
     """
     if use_compat not in [None, True, False]:
@@ -297,6 +348,15 @@ def array_namespace(*xs, api_version=None, use_compat=None):
                 # not have a wrapper submodule for it.
                 import jax.experimental.array_api as jnp
             namespaces.add(jnp)
+        elif is_pydata_sparse_array(x):
+            if use_compat is True:
+                _check_api_version(api_version)
+                raise ValueError("`sparse` does not have an array-api-compat wrapper")
+            else:
+                import sparse
+            # `sparse` is already an array namespace. We do not have a wrapper
+            # submodule for it.
+            namespaces.add(sparse)
         elif hasattr(x, '__array_namespace__'):
             if use_compat is True:
                 raise ValueError("The given array does not have an array-api-compat wrapper")
@@ -391,8 +451,23 @@ def device(x: Array, /) -> Device:
             return x.device()
         else:
             return x.device
+    elif is_pydata_sparse_array(x):
+        # `sparse` will gain `.device`, so check for this first.
+        x_device = getattr(x, 'device', None)
+        if x_device is not None:
+            return x_device
+        # Everything but DOK has this attr.
+        try:
+            inner = x.data
+        except AttributeError:
+            return "cpu"
+        # Return the device of the constituent array
+        return device(inner)
     return x.device
 
+# Prevent shadowing, used below
+_device = device
+
 # Based on cupy.array_api.Array.to_device
 def _cupy_to_device(x, device, /, stream=None):
     import cupy as cp
@@ -508,6 +583,10 @@ def to_device(x: Array, device: Device, /, *, stream: Optional[Union[int, Any]]
         # This import adds to_device to x
         import jax.experimental.array_api # noqa: F401
         return x.to_device(device, stream=stream)
+    elif is_pydata_sparse_array(x) and device == _device(x):
+        # Perform trivial check to return the same array if
+        # device is same instead of err-ing.
+        return x
     return x.to_device(device, stream=stream)
 
 def size(x):
@@ -534,6 +613,7 @@ __all__ = [
     "is_jax_array",
     "is_numpy_array",
     "is_torch_array",
+    "is_pydata_sparse_array",
     "size",
     "to_device",
 ]

{array_api_compat-1.6 → array_api_compat-1.7}/array_api_compat/torch/linalg.py

@@ -60,6 +60,22 @@ def vecdot(x1: array, x2: array, /, *, axis: int = -1, **kwargs) -> array:
 
 def solve(x1: array, x2: array, /, **kwargs) -> array:
     x1, x2 = _fix_promotion(x1, x2, only_scalar=False)
+    # Torch tries to emulate NumPy 1 solve behavior by using batched 1-D solve
+    # whenever
+    # 1. x1.ndim - 1 == x2.ndim
+    # 2. x1.shape[:-1] == x2.shape
+    #
+    # See linalg_solve_is_vector_rhs in
+    # aten/src/ATen/native/LinearAlgebraUtils.h and
+    # TORCH_META_FUNC(_linalg_solve_ex) in
+    # aten/src/ATen/native/BatchLinearAlgebra.cpp in the PyTorch source code.
+    #
+    # The easiest way to work around this is to prepend a size 1 dimension to
+    # x2, since x2 is already one dimension less than x1.
+    #
+    # See https://github.com/pytorch/pytorch/issues/52915
+    if x2.ndim != 1 and x1.ndim - 1 == x2.ndim and x1.shape[:-1] == x2.shape:
+        x2 = x2[None]
     return torch.linalg.solve(x1, x2, **kwargs)
 
 # torch.trace doesn't support the offset argument and doesn't support stacking
@@ -78,7 +94,23 @@ def vector_norm(
 ) -> array:
     # torch.vector_norm incorrectly treats axis=() the same as axis=None
     if axis == ():
-        keepdims = True
+        out = kwargs.get('out')
+        if out is None:
+            dtype = None
+            if x.dtype == torch.complex64:
+                dtype = torch.float32
+            elif x.dtype == torch.complex128:
+                dtype = torch.float64
+
+            out = torch.zeros_like(x, dtype=dtype)
+
+        # The norm of a single scalar works out to abs(x) in every case except
+        # for ord=0, which is x != 0.
+        if ord == 0:
+            out[:] = (x != 0)
+        else:
+            out[:] = torch.abs(x)
+        return out
     return torch.linalg.vector_norm(x, ord=ord, axis=axis, keepdim=keepdims, **kwargs)
 
 __all__ = linalg_all + ['outer', 'matmul', 'matrix_transpose', 'tensordot',

{array_api_compat-1.6 → array_api_compat-1.7}/array_api_compat.egg-info/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.1
 Name: array_api_compat
-Version: 1.6
+Version: 1.7
 Summary: A wrapper around NumPy and other array libraries to make them compatible with the Array API standard
 Home-page: https://data-apis.org/array-api-compat/
 Author: Consortium for Python Data API Standards
@@ -23,13 +23,15 @@ Provides-Extra: pytorch
 Requires-Dist: pytorch; extra == "pytorch"
 Provides-Extra: dask
 Requires-Dist: dask; extra == "dask"
+Provides-Extra: sprase
+Requires-Dist: sparse>=0.15.1; extra == "sprase"
 
 # Array API compatibility library
 
 This is a small wrapper around common array libraries that is compatible with
 the [Array API standard](https://data-apis.org/array-api/latest/). Currently,
-NumPy, CuPy, PyTorch, Dask, and JAX are supported. If you want support for other array
-libraries, or if you encounter any issues, please [open an
+NumPy, CuPy, PyTorch, Dask, JAX and `sparse` are supported. If you want support
+for other array libraries, or if you encounter any issues, please [open an
 issue](https://github.com/data-apis/array-api-compat/issues).
 
 See the documentation for more details https://data-apis.org/array-api-compat/

{array_api_compat-1.6 → array_api_compat-1.7}/array_api_compat.egg-info/requires.txt

@@ -13,3 +13,6 @@ numpy
 
 [pytorch]
 pytorch
+
+[sprase]
+sparse>=0.15.1

{array_api_compat-1.6 → array_api_compat-1.7}/setup.py

@@ -21,6 +21,7 @@ setup(
         "jax": "jax",
         "pytorch": "pytorch",
         "dask": "dask",
+        "sprase": "sparse >=0.15.1",
     },
     classifiers=[
         "Programming Language :: Python :: 3",

{array_api_compat-1.6 → array_api_compat-1.7}/tests/test_array_namespace.py

@@ -2,6 +2,7 @@ import subprocess
 import sys
 import warnings
 
+import jax
 import numpy as np
 import pytest
 import torch
@@ -18,7 +19,7 @@ def test_array_namespace(library, api_version, use_compat):
     xp = import_(library)
 
     array = xp.asarray([1.0, 2.0, 3.0])
-    if use_compat is True and library in ['array_api_strict', 'jax.numpy']:
+    if use_compat is True and library in {'array_api_strict', 'jax.numpy', 'sparse'}:
         pytest.raises(ValueError, lambda: array_namespace(array, use_compat=use_compat))
         return
     namespace = array_api_compat.array_namespace(array, api_version=api_version, use_compat=use_compat)
@@ -55,6 +56,12 @@ assert namespace == jax.experimental.array_api
 """
         subprocess.run([sys.executable, "-c", code], check=True)
 
+def test_jax_zero_gradient():
+    jx = jax.numpy.arange(4)
+    jax_zero = jax.vmap(jax.grad(jax.numpy.float32, allow_int=True))(jx)
+    assert (array_api_compat.get_namespace(jax_zero) is
+            array_api_compat.get_namespace(jx))
+
 def test_array_namespace_errors():
     pytest.raises(TypeError, lambda: array_namespace([1]))
     pytest.raises(TypeError, lambda: array_namespace())

{array_api_compat-1.6 → array_api_compat-1.7}/tests/test_common.py

@@ -1,5 +1,5 @@
 from array_api_compat import (is_numpy_array, is_cupy_array, is_torch_array, # noqa: F401
-                              is_dask_array, is_jax_array)
+                              is_dask_array, is_jax_array, is_pydata_sparse_array)
 
 from array_api_compat import is_array_api_obj, device, to_device
 
@@ -16,6 +16,7 @@ is_functions = {
     'torch': 'is_torch_array',
     'dask.array': 'is_dask_array',
     'jax.numpy': 'is_jax_array',
+    'sparse': 'is_pydata_sparse_array',
 }
 
 @pytest.mark.parametrize('library', is_functions.keys())
@@ -76,6 +77,8 @@ def test_asarray_cross_library(source_library, target_library, request):
     if source_library == "cupy" and target_library != "cupy":
         # cupy explicitly disallows implicit conversions to CPU
         pytest.skip(reason="cupy does not support implicit conversion to CPU")
+    elif source_library == "sparse" and target_library != "sparse":
+        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]]

{array_api_compat-1.6 → array_api_compat-1.7}/tests/test_no_dependencies.py

@@ -33,7 +33,7 @@ def _test_dependency(mod):
 
     # array-api-strict is an example of an array API library that isn't
     # wrapped by array-api-compat.
-    if "strict" not in mod:
+    if "strict" not in mod and mod != "sparse":
         is_mod_array = getattr(array_api_compat, f"is_{mod.split('.')[0]}_array")
         assert not is_mod_array(a)
         assert mod not in sys.modules
@@ -50,7 +50,7 @@ def _test_dependency(mod):
 # Y (except most array libraries actually do themselves depend on numpy).
 
 @pytest.mark.parametrize("library", ["cupy", "numpy", "torch", "dask.array",
-                                     "jax.numpy", "array_api_strict"])
+                                     "jax.numpy", "sparse", "array_api_strict"])
 def test_numpy_dependency(library):
     # This import is here because it imports numpy
     from ._helpers import import_