Merge branch 'main' into np-solve-fix

Author: asmeurer

.github/workflows/ruff.yml

@@ -16,4 +16,4 @@ jobs:
           pip install ruff
       # Update output format to enable automatic inline annotations.
       - name: Run Ruff
-        run: ruff check --output-format=github --select F822,PLC0414,RUF022 --preview .
+        run: ruff check --output-format=github .

.gitignore

@@ -127,3 +127,6 @@ dmypy.json
 
 # Pyre type checker
 .pyre/
+
+# macOS specific iles
+.DS_Store

README.md

@@ -71,23 +71,24 @@ namespace, except that functions that are part of the array API are wrapped so
 that they have the correct array API behavior. In each case, the array object
 used will be the same array object from the wrapped library.
 
-## Difference between `array_api_compat` and `numpy.array_api`
+## Difference between `array_api_compat` and `array_api_strict`
 
-`numpy.array_api` is a strict minimal implementation of the Array API (see
+`array_api_strict` is a strict minimal implementation of the array API standard, formerly
+known as `numpy.array_api` (see
 [NEP 47](https://numpy.org/neps/nep-0047-array-api-standard.html)). For
-example, `numpy.array_api` does not include any functions that are not part of
+example, `array_api_strict` does not include any functions that are not part of
 the array API specification, and will explicitly disallow behaviors that are
 not required by the spec (e.g., [cross-kind type
 promotions](https://data-apis.org/array-api/latest/API_specification/type_promotion.html)).
-(`cupy.array_api` is similar to `numpy.array_api`)
+(`cupy.array_api` is similar to `array_api_strict`)
 
 `array_api_compat`, on the other hand, is just an extension of the
 corresponding array library namespaces with changes needed to be compliant
 with the array API. It includes all additional library functions not mentioned
 in the spec, and allows any library behaviors not explicitly disallowed by it,
 such as cross-kind casting.
 
-In particular, unlike `numpy.array_api`, this package does not use a separate
+In particular, unlike `array_api_strict`, this package does not use a separate
 `Array` object, but rather just uses the corresponding array library array
 objects (`numpy.ndarray`, `cupy.ndarray`, `torch.Tensor`, etc.) directly. This
 is because those are the objects that are going to be passed as inputs to
@@ -96,7 +97,7 @@ functions by end users. This does mean that a few behaviors cannot be wrapped
 most things.
 
 Array consuming library authors coding against the array API may wish to test
-against `numpy.array_api` to ensure they are not using functionality outside
+against `array_api_strict` to ensure they are not using functionality outside
 of the standard, but prefer this implementation for the default behavior for
 end-users.
 
@@ -125,11 +126,11 @@ part of the specification but which are useful for using the array API:
   [`x.device`](https://data-apis.org/array-api/latest/API_specification/generated/signatures.array_object.array.device.html)
   in the array API specification. Included because `numpy.ndarray` does not
   include the `device` attribute and this library does not wrap or extend the
-  array object. Note that for NumPy, `device(x)` is always `"cpu"`.
+  array object. Note that for NumPy and dask, `device(x)` is always `"cpu"`.
 
 - `to_device(x, device, /, *, stream=None)`: Equivalent to
   [`x.to_device`](https://data-apis.org/array-api/latest/API_specification/generated/signatures.array_object.array.to_device.html).
-  Included because neither NumPy's, CuPy's, nor PyTorch's array objects
+  Included because neither NumPy's, CuPy's, Dask's, nor PyTorch's array objects
   include this method. For NumPy, this function effectively does nothing since
   the only supported device is the CPU, but for CuPy, this method supports
   CuPy CUDA
@@ -240,6 +241,30 @@ Unlike the other libraries supported here, JAX array API support is contained
 entirely in the JAX library. The JAX array API support is tracked at
 https://github.com/google/jax/issues/18353.
 
+## Dask
+
+If you're using dask with numpy, many of the same limitations that apply to numpy
+will also apply to dask. Besides those differences, other limitations include missing
+sort functionality (no `sort` or `argsort`), and limited support for the optional `linalg`
+and `fft` extensions.
+
+In particular, the `fft` namespace is not compliant with the array API spec. Any functions
+that you find under the `fft` namespace are the original, unwrapped functions under [`dask.array.fft`](https://docs.dask.org/en/latest/array-api.html#fast-fourier-transforms), which may or may not be Array API compliant. Use at your own risk!
+
+For `linalg`, several methods are missing, for example:
+- `cross`
+- `det`
+- `eigh`
+- `eigvalsh`
+- `matrix_power`
+- `pinv`
+- `slogdet`
+- `matrix_norm`
+- `matrix_rank`
+Other methods may only be partially implemented or return incorrect results at times.
+
+The minimum supported Dask version is 2023.12.0.
+
 ## Vendoring
 
 This library supports vendoring as an installation method. To vendor the

array_api_compat/__init__.py

@@ -5,14 +5,14 @@
 Array API standard https://data-apis.org/array-api/latest/. See also NEP 47
 https://numpy.org/neps/nep-0047-array-api-standard.html.
 
-Unlike numpy.array_api, this is not a strict minimal implementation of the
+Unlike array_api_strict, this is not a strict minimal implementation of the
 Array API, but rather just an extension of the main NumPy namespace with
 changes needed to be compliant with the Array API. See
 https://numpy.org/doc/stable/reference/array_api.html for a full list of
-changes. In particular, unlike numpy.array_api, this package does not use a
+changes. In particular, unlike array_api_strict, this package does not use a
 separate Array object, but rather just uses numpy.ndarray directly.
 
-Library authors using the Array API may wish to test against numpy.array_api
+Library authors using the Array API may wish to test against array_api_strict
 to ensure they are not using functionality outside of the standard, but prefer
 this implementation for the default when working with NumPy arrays.
 

array_api_compat/_internal.py

@@ -5,7 +5,6 @@
 from functools import wraps
 from inspect import signature
 
-
 def get_xp(xp):
     """
     Decorator to automatically replace xp with the corresponding array module.
@@ -45,31 +44,3 @@ def wrapped_f(*args, **kwargs):
         return wrapped_f
 
     return inner
-
-
-def _get_all_public_members(module, exclude=None, extend_all=False):
-    """Get all public members of a module.
-
-    Parameters
-    ----------
-    module : module
-        The module to get members from.
-    exclude : callable, optional
-        A callable that takes a name and returns True if the name should be
-        excluded from the list of members.
-    extend_all : bool, optional
-        If True, extend the module's __all__ attribute with the members of the
-        module derived from dir(module). To be used for libraries that do not have a complete __all__ list.
-    """
-    members = getattr(module, "__all__", [])
-
-    if members and not extend_all:
-        return members
-
-    if exclude is None:
-        exclude = lambda name: name.startswith("_")  # noqa: E731
-
-    members = members + [_ for _ in dir(module) if not exclude(_)]
-
-    # remove duplicates
-    return list(set(members))

array_api_compat/common/__init__.py

@@ -1,27 +1 @@
-from ._helpers import (
-    array_namespace,
-    device,
-    get_namespace,
-    is_array_api_obj,
-    is_cupy_array,
-    is_dask_array,
-    is_jax_array,
-    is_numpy_array,
-    is_torch_array,
-    size,
-    to_device,
-)
-
-__all__ = [
-    "array_namespace",
-    "device",
-    "get_namespace",
-    "is_array_api_obj",
-    "is_cupy_array",
-    "is_dask_array",
-    "is_jax_array",
-    "is_numpy_array",
-    "is_torch_array",
-    "size",
-    "to_device",
-]
+from ._helpers import * # noqa: F403

array_api_compat/common/_aliases.py

@@ -146,6 +146,9 @@ def zeros_like(
 
 # The functions here return namedtuples (np.unique() returns a normal
 # tuple).
+
+# Note that these named tuples aren't actually part of the standard namespace,
+# but I don't see any issue with exporting the names here regardless.
 class UniqueAllResult(NamedTuple):
     values: ndarray
     indices: ndarray
@@ -543,5 +546,13 @@ def isdtype(
         # This will allow things that aren't required by the spec, like
         # isdtype(np.float64, float) or isdtype(np.int64, 'l'). Should we be
         # more strict here to match the type annotation? Note that the
-        # numpy.array_api implementation will be very strict.
+        # array_api_strict implementation will be very strict.
         return dtype == kind
+
+__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', 'permute_dims', 'reshape', 'argsort',
+           'sort', 'nonzero', 'sum', 'prod', 'ceil', 'floor', 'trunc',
+           'matmul', 'matrix_transpose', 'tensordot', 'vecdot', 'isdtype']

array_api_compat/common/_helpers.py

@@ -159,7 +159,16 @@ def _check_device(xp, device):
         if device not in ["cpu", None]:
             raise ValueError(f"Unsupported device for NumPy: {device!r}")
 
-# device() is not on numpy.ndarray and to_device() is not on numpy.ndarray
+# Placeholder object to represent the dask device
+# when the array backend is not the CPU.
+# (since it is not easy to tell which device a dask array is on)
+class _dask_device:
+    def __repr__(self):
+        return "DASK_DEVICE"
+
+_DASK_DEVICE = _dask_device()
+
+# device() is not on numpy.ndarray or dask.array and to_device() is not on numpy.ndarray
 # or cupy.ndarray. They are not included in array objects of this library
 # because this library just reuses the respective ndarray classes without
 # wrapping or subclassing them. These helper functions can be used instead of
@@ -181,7 +190,17 @@ def device(x: Array, /) -> Device:
     """
     if is_numpy_array(x):
         return "cpu"
-    if is_jax_array(x):
+    elif is_dask_array(x):
+        # Peek at the metadata of the jax array to determine type
+        try:
+            import numpy as np
+            if isinstance(x._meta, np.ndarray):
+                # Must be on CPU since backed by numpy
+                return "cpu"
+        except ImportError:
+            pass
+        return _DASK_DEVICE
+    elif is_jax_array(x):
         # JAX has .device() as a method, but it is being deprecated so that it
         # can become a property, in accordance with the standard. In order for
         # this function to not break when JAX makes the flip, we check for
@@ -288,3 +307,19 @@ def size(x):
     if None in x.shape:
         return None
     return math.prod(x.shape)
+
+__all__ = [
+    "array_namespace",
+    "device",
+    "get_namespace",
+    "is_array_api_obj",
+    "is_cupy_array",
+    "is_dask_array",
+    "is_jax_array",
+    "is_numpy_array",
+    "is_torch_array",
+    "size",
+    "to_device",
+]
+
+_all_ignore = ['sys', 'math', 'inspect']

array_api_compat/common/_linalg.py

@@ -5,13 +5,15 @@
     from typing import Literal, Optional, Tuple, Union
     from ._typing import ndarray
 
+import math
+
 import numpy as np
 if np.__version__[0] == "2":
     from numpy.lib.array_utils import normalize_axis_tuple
 else:
     from numpy.core.numeric import normalize_axis_tuple
 
-from ._aliases import matrix_transpose, isdtype
+from ._aliases import matmul, matrix_transpose, tensordot, vecdot, isdtype
 from .._internal import get_xp
 
 # These are in the main NumPy namespace but not in numpy.linalg
@@ -110,21 +112,22 @@ def vector_norm(x: ndarray, /, xp, *, axis: Optional[Union[int, Tuple[int, ...]]
     # on a single dimension.
     if axis is None:
         # Note: xp.linalg.norm() doesn't handle 0-D arrays
-        x = x.ravel()
+        _x = x.ravel()
         _axis = 0
     elif isinstance(axis, tuple):
         # Note: The axis argument supports any number of axes, whereas
         # xp.linalg.norm() only supports a single axis for vector norm.
         normalized_axis = normalize_axis_tuple(axis, x.ndim)
         rest = tuple(i for i in range(x.ndim) if i not in normalized_axis)
         newshape = axis + rest
-        x = xp.transpose(x, newshape).reshape(
-            (xp.prod([x.shape[i] for i in axis], dtype=int), *[x.shape[i] for i in rest]))
+        _x = xp.transpose(x, newshape).reshape(
+            (math.prod([x.shape[i] for i in axis]), *[x.shape[i] for i in rest]))
         _axis = 0
     else:
+        _x = x
         _axis = axis
 
-    res = xp.linalg.norm(x, axis=_axis, ord=ord)
+    res = xp.linalg.norm(_x, axis=_axis, ord=ord)
 
     if keepdims:
         # We can't reuse xp.linalg.norm(keepdims) because of the reshape hacks
@@ -150,3 +153,9 @@ def trace(x: ndarray, /, xp, *, offset: int = 0, dtype=None, **kwargs) -> ndarra
         elif x.dtype == xp.complex64:
             dtype = xp.complex128
     return xp.asarray(xp.trace(x, offset=offset, dtype=dtype, axis1=-2, axis2=-1, **kwargs))
+
+__all__ = ['cross', 'matmul', 'outer', 'tensordot', 'EighResult',
+           'QRResult', 'SlogdetResult', 'SVDResult', 'eigh', 'qr', 'slogdet',
+           'svd', 'cholesky', 'matrix_rank', 'pinv', 'matrix_norm',
+           'matrix_transpose', 'svdvals', 'vecdot', 'vector_norm', 'diagonal',
+           'trace']

array_api_compat/common/_typing.py

@@ -20,4 +20,4 @@ def __len__(self, /) -> int: ...
 SupportsBufferProtocol = Any
 
 Array = Any
-Device = Any
+Device = Any