Merge branch 'add-dask' of github.com:lithomas1/array-api-compat into add-dask

Author: lithomas1

.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

@@ -30,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:
@@ -51,19 +51,20 @@ jobs:
         submodules: 'true'
         path: array-api-tests
     - name: Set up Python ${{ matrix.python-version }}
-      uses: actions/setup-python@v4
+      uses: actions/setup-python@v5
       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 }}' ${{ 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.module-name || inputs.package-name }}
         # This enables the NEP 50 type promotion behavior (without it a lot of

.github/workflows/publish-package.yml

@@ -35,7 +35,7 @@ jobs:
         fetch-depth: 0
 
     - name: Set up Python
-      uses: actions/setup-python@v4
+      uses: actions/setup-python@v5
       with:
         python-version: '3.x'
 
@@ -59,7 +59,7 @@ jobs:
       run: python -m zipfile --list dist/array_api_compat-*.whl
 
     - name: Upload distribution artifact
-      uses: actions/upload-artifact@v3
+      uses: actions/upload-artifact@v4
       with:
         name: dist-artifact
         path: dist
@@ -80,7 +80,7 @@ jobs:
 
     steps:
     - name: Download distribution artifact
-      uses: actions/download-artifact@v3
+      uses: actions/download-artifact@v4
       with:
         name: dist-artifact
         path: dist

.github/workflows/tests.yml

@@ -9,7 +9,7 @@ jobs:
       fail-fast: true
     steps:
       - uses: actions/checkout@v4
-      - uses: actions/setup-python@v4
+      - uses: actions/setup-python@v5
         with:
           python-version: ${{ matrix.python-version }}
       - name: Install Dependencies

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

@@ -300,3 +300,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

@@ -17,6 +17,6 @@
 this implementation for the default when working with NumPy arrays.
 
 """
-__version__ = '1.4'
+__version__ = '1.4.1'
 
 from .common import *

array_api_compat/common/_aliases.py

@@ -396,6 +396,12 @@ def sort(
         res = xp.flip(res, axis=axis)
     return res
 
+# nonzero should error for zero-dimensional arrays
+def nonzero(x: ndarray, /, xp, **kwargs) -> Tuple[ndarray, ...]:
+    if x.ndim == 0:
+        raise ValueError("nonzero() does not support zero-dimensional arrays")
+    return xp.nonzero(x, **kwargs)
+
 # sum() and prod() should always upcast when dtype=None
 def sum(
     x: ndarray,
@@ -536,5 +542,5 @@ def isdtype(
            'UniqueAllResult', 'UniqueCountsResult', 'UniqueInverseResult',
            'unique_all', 'unique_counts', 'unique_inverse', 'unique_values',
            'astype', 'std', 'var', 'permute_dims', 'reshape', 'argsort',
-           'sort', 'sum', 'prod', 'ceil', 'floor', 'trunc', 'matmul',
-           'matrix_transpose', 'tensordot', 'vecdot', 'isdtype']
+           'sort', 'nonzero', 'sum', 'prod', 'ceil', 'floor', 'trunc',
+           'matmul', 'matrix_transpose', 'tensordot', 'vecdot', 'isdtype']

array_api_compat/common/_helpers.py

@@ -81,9 +81,7 @@ def your_function(x, y):
     """
     namespaces = set()
     for x in xs:
-        if hasattr(x, '__array_namespace__'):
-            namespaces.add(x.__array_namespace__(api_version=api_version))
-        elif _is_numpy_array(x):
+        if _is_numpy_array(x):
             _check_api_version(api_version)
             if _use_compat:
                 from .. import numpy as numpy_namespace
@@ -114,9 +112,11 @@ def your_function(x, y):
                 namespaces.add(dask_namespace)
             else:
                 raise TypeError("_use_compat cannot be False if input array is a dask array!")
+        elif hasattr(x, '__array_namespace__'):
+            namespaces.add(x.__array_namespace__(api_version=api_version))
         else:
             # TODO: Support Python scalars?
-            raise TypeError("The input is not a supported array type")
+            raise TypeError(f"{type(x).__name__} is not a supported array type")
 
     if not namespaces:
         raise TypeError("Unrecognized array input")

array_api_compat/common/_linalg.py

@@ -11,7 +11,7 @@
 else:
     from numpy.core.numeric import normalize_axis_tuple
 
-from ._aliases import matmul, matrix_transpose, tensordot, vecdot
+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
@@ -59,7 +59,10 @@ def svd(x: ndarray, /, xp, *, full_matrices: bool = True, **kwargs) -> SVDResult
 def cholesky(x: ndarray, /, xp, *, upper: bool = False, **kwargs) -> ndarray:
     L = xp.linalg.cholesky(x, **kwargs)
     if upper:
-        return get_xp(xp)(matrix_transpose)(L)
+        U = get_xp(xp)(matrix_transpose)(L)
+        if get_xp(xp)(isdtype)(U.dtype, 'complex floating'):
+            U = xp.conj(U)
+        return U
     return L
 
 # The rtol keyword argument of matrix_rank() and pinv() is new from NumPy.