Merge branch 'main' into fft

Author: asmeurer

.github/workflows/array-api-tests-dask.yml

@@ -0,0 +1,12 @@
+name: Array API Tests (Dask)
+
+on: [push, pull_request]
+
+jobs:
+  array-api-tests-dask:
+    uses: ./.github/workflows/array-api-tests.yml
+    with:
+      package-name: dask
+      module-name: dask.array
+      extra-requires: numpy
+      pytest-extra-args: --disable-deadline --max-examples=5

.github/workflows/array-api-tests-numpy-dev.yml

@@ -0,0 +1,11 @@
+name: Array API Tests (NumPy dev)
+
+on: [push, pull_request]
+
+jobs:
+  array-api-tests-numpy-dev:
+    uses: ./.github/workflows/array-api-tests.yml
+    with:
+      package-name: numpy
+      extra-requires: '--pre --extra-index https://pypi.anaconda.org/scientific-python-nightly-wheels/simple'
+      xfails-file-extra: '-dev'

.github/workflows/array-api-tests.yml

@@ -6,6 +6,12 @@ on:
       package-name:
         required: true
         type: string
+      module-name:
+        required: false
+        type: string
+      extra-requires:
+        required: false
+        type: string
       package-version:
         required: false
         type: string
@@ -24,7 +30,7 @@ on:
 
 
 env:
-  PYTEST_ARGS: "--max-examples 200 -v -rxXfE --ci ${{ inputs.pytest-extra-args }}"
+  PYTEST_ARGS: "--max-examples 200 -v -rxXfE --ci ${{ inputs.pytest-extra-args }} --hypothesis-disable-deadline"
 
 jobs:
   tests:
@@ -49,17 +55,18 @@ jobs:
       with:
         python-version: ${{ matrix.python-version }}
     - name: Install dependencies
-      # NumPy 1.21 doesn't support Python 3.11. There doesn't seem to be a way
-      # to put this in the numpy 1.21 config file.
-      if: "! (matrix.python-version == '3.11' && inputs.package-name == 'numpy' && contains(inputs.package-version, '1.21'))"
+      # NumPy 1.21 doesn't support Python 3.11. NumPy 2.0 doesn't support
+      # Python 3.8. There doesn't seem to be a way to put this in the numpy
+      # 1.21 config file.
+      if: "! ((matrix.python-version == '3.11' && inputs.package-name == 'numpy' && contains(inputs.package-version, '1.21')) || (matrix.python-version == '3.8' && inputs.package-name == 'numpy' && contains(inputs.xfails-file-extra, 'dev')))"
       run: |
         python -m pip install --upgrade pip
-        python -m pip install '${{ inputs.package-name }} ${{ inputs.package-version }}'
+        python -m pip install '${{ inputs.package-name }} ${{ inputs.package-version }}' ${{ inputs.extra-requires }}
         python -m pip install -r ${GITHUB_WORKSPACE}/array-api-tests/requirements.txt
     - name: Run the array API testsuite (${{ inputs.package-name }})
-      if: "! (matrix.python-version == '3.11' && inputs.package-name == 'numpy' && contains(inputs.package-version, '1.21'))"
+      if: "! ((matrix.python-version == '3.11' && inputs.package-name == 'numpy' && contains(inputs.package-version, '1.21')) || (matrix.python-version == '3.8' && inputs.package-name == 'numpy' && contains(inputs.xfails-file-extra, 'dev')))"
       env:
-        ARRAY_API_TESTS_MODULE: array_api_compat.${{ inputs.package-name }}
+        ARRAY_API_TESTS_MODULE: array_api_compat.${{ inputs.module-name || inputs.package-name }}
         # This enables the NEP 50 type promotion behavior (without it a lot of
         # tests fail on bad scalar type promotion behavior)
         NPY_PROMOTION_STATE: weak

.github/workflows/ruff.yml

@@ -0,0 +1,19 @@
+name: CI
+on: [push, pull_request]
+jobs:
+  check-ruff:
+    runs-on: ubuntu-latest
+    continue-on-error: true
+    steps:
+      - uses: actions/checkout@v4
+      - name: Install Python
+        uses: actions/setup-python@v5
+        with:
+          python-version: "3.11"
+      - name: Install dependencies
+        run: |
+          python -m pip install --upgrade pip
+          pip install ruff
+      # Update output format to enable automatic inline annotations.
+      - name: Run Ruff
+        run: ruff check --output-format=github .

.github/workflows/tests.yml

@@ -15,7 +15,7 @@ jobs:
       - name: Install Dependencies
         run: |
           python -m pip install --upgrade pip
-          python -m pip install pytest numpy torch
+          python -m pip install pytest numpy torch dask[array] jax[cpu]
 
       - name: Run Tests
         run: |

.gitignore

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

CHANGELOG.md

@@ -1,3 +1,22 @@
+# 1.4.1 (2024-01-18)
+
+## Minor Changes
+
+- Add support for the upcoming NumPy 2.0 release.
+
+- Added a torch wrapper for `trace` (`torch.trace` doesn't support the
+  `offset` argument or stacking)
+
+- Wrap numpy, cupy, and torch `nonzero` to raise an error for zero-dimensional
+  input arrays.
+
+- Add torch wrapper for `newaxis`.
+
+- Improve error message for `array_namespace`
+
+- Fix linalg.cholesky returning the conjugate of the expected upper
+  decomposition for numpy and cupy.
+
 # 1.4 (2023-09-13)
 
 ## Major Changes

README.md

@@ -2,7 +2,7 @@
 
 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, and PyTorch are supported. If you want support for other array
+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
 issue](https://github.com/data-apis/array-api-compat/issues).
 
@@ -56,28 +56,39 @@ import array_api_compat.cupy as cp
 import array_api_compat.torch as torch
 ```
 
-Each will include all the functions from the normal NumPy/CuPy/PyTorch
+```py
+import array_api_compat.dask as da
+```
+
+> [!NOTE]
+> There is no `array_api_compat.jax` submodule. JAX support is contained
+> in JAX itself in the `jax.experimental.array_api` module. array-api-compat simply
+> wraps that submodule. The main JAX support in this module consists of
+> supporting it in the [helper functions](#helper-functions) defined below.
+
+Each will include all the functions from the normal NumPy/CuPy/PyTorch/dask.array
 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
@@ -86,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.
 
@@ -99,6 +110,11 @@ part of the specification but which are useful for using the array API:
 - `is_array_api_obj(x)`: Return `True` if `x` is an array API compatible array
   object.
 
+- `is_numpy_array(x)`, `is_cupy_array(x)`, `is_torch_array(x)`,
+  `is_dask_array(x)`, `is_jax_array(x)`: return `True` if `x` is an array from
+  the corresponding library. These functions do not import the underlying
+  library if it has not already been imported, so they are cheap to use.
+
 - `array_namespace(*xs)`: Get the corresponding array API namespace for the
   arrays `xs`. For example, if the arrays are NumPy arrays, the returned
   namespace will be `array_api_compat.numpy`. Note that this function will
@@ -110,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
@@ -219,6 +235,36 @@ version.
 
 The minimum supported PyTorch version is 1.13.
 
+### JAX
+
+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
@@ -300,3 +346,54 @@ corresponding document does not yet exist for PyTorch, but you can examine the
 various comments in the
 [implementation](https://github.com/data-apis/array-api-compat/blob/main/array_api_compat/torch/_aliases.py)
 to see what functions and behaviors have been wrapped.
+
+
+## Releasing
+
+To release, first note that CuPy must be tested manually (it isn't tested on
+CI). Use the script
+
+```
+./test_cupy.sh
+```
+
+on a machine with a CUDA GPU.
+
+Once you are ready to release, create a PR with a release branch, so that you
+can verify that CI is passing. You must edit
+
+```
+array_api_compat/__init__.py
+```
+
+and update the version (the version is not computed from the tag because that
+would break vendorability). You should also edit
+
+```
+CHANGELOG.md
+```
+
+with the changes for the release.
+
+Then create a tag
+
+```
+git tag -a <version>
+```
+
+and push it to GitHub
+
+```
+git push origin <version>
+```
+
+Check that the `publish distributions` action works. Note that this action
+will run even if the other CI fails, so you must make sure that CI is passing
+*before* tagging.
+
+This does mean you can ignore CI failures, but ideally you should fix any
+failures or update the `*-xfails.txt` files before tagging, so that CI and the
+cupy tests pass. Otherwise it will be hard to tell what things are breaking in
+the future. It's also a good idea to remove any xpasses from those files (but
+be aware that some xfails are from flaky failures, so unless you know the
+underlying issue has been fixed, a xpass test is probably still xfail).

array_api_compat/__init__.py

@@ -5,18 +5,18 @@
 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.
 
 """
-__version__ = '1.4'
+__version__ = '1.4.1'
 
-from .common import *
+from .common import *  # noqa: F401, F403

array_api_compat/_internal.py

@@ -21,13 +21,16 @@ def func(x, /, xp, kwarg=None):
     arguments.
 
     """
+
     def inner(f):
         @wraps(f)
         def wrapped_f(*args, **kwargs):
             return f(*args, xp=xp, **kwargs)
 
         sig = signature(f)
-        new_sig = sig.replace(parameters=[sig.parameters[i] for i in sig.parameters if i != 'xp'])
+        new_sig = sig.replace(
+            parameters=[sig.parameters[i] for i in sig.parameters if i != "xp"]
+        )
 
         if wrapped_f.__doc__ is None:
             wrapped_f.__doc__ = f"""\