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

lithomas1 · lithomas1 · commit 565666aa1a22 · 2024-01-24T11:14:09.000-05:00
diff --git a/.github/workflows/array-api-tests-numpy-dev.yml b/.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'
diff --git a/.github/workflows/array-api-tests.yml b/.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
diff --git a/.github/workflows/publish-package.yml b/.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
diff --git a/.github/workflows/tests.yml b/.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
diff --git a/CHANGELOG.md b/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
diff --git a/README.md b/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).
diff --git a/array_api_compat/__init__.py b/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 *
diff --git a/array_api_compat/common/_aliases.py b/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']
diff --git a/array_api_compat/common/_helpers.py b/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")
diff --git a/array_api_compat/common/_linalg.py b/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.
diff --git a/array_api_compat/cupy/_aliases.py b/array_api_compat/cupy/_aliases.py
@@ -52,6 +52,7 @@
 reshape = get_xp(cp)(_aliases.reshape)
 argsort = get_xp(cp)(_aliases.argsort)
 sort = get_xp(cp)(_aliases.sort)
+nonzero = get_xp(cp)(_aliases.nonzero)
 sum = get_xp(cp)(_aliases.sum)
 prod = get_xp(cp)(_aliases.prod)
 ceil = get_xp(cp)(_aliases.ceil)
@@ -60,8 +61,17 @@
 matmul = get_xp(cp)(_aliases.matmul)
 matrix_transpose = get_xp(cp)(_aliases.matrix_transpose)
 tensordot = get_xp(cp)(_aliases.tensordot)
-vecdot = get_xp(cp)(_aliases.vecdot)
-isdtype = get_xp(cp)(_aliases.isdtype)
+
+# 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'):
+    vecdot = cp.vecdot
+else:
+    vecdot = get_xp(cp)(_aliases.vecdot)
+if hasattr(cp, 'isdtype'):
+    isdtype = cp.isdtype
+else:
+    isdtype = get_xp(cp)(_aliases.isdtype)
 
 __all__ = _aliases.__all__ + ['asarray', 'asarray_cupy', 'bool', 'acos',
                               'acosh', 'asin', 'asinh', 'atan', 'atan2',
diff --git a/array_api_compat/cupy/linalg.py b/array_api_compat/cupy/linalg.py
@@ -29,10 +29,16 @@
 pinv = get_xp(cp)(_linalg.pinv)
 matrix_norm = get_xp(cp)(_linalg.matrix_norm)
 svdvals = get_xp(cp)(_linalg.svdvals)
-vector_norm = get_xp(cp)(_linalg.vector_norm)
 diagonal = get_xp(cp)(_linalg.diagonal)
 trace = get_xp(cp)(_linalg.trace)
 
+# 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.linalg, 'vector_norm'):
+    vector_norm = cp.linalg.vector_norm
+else:
+    vector_norm = get_xp(cp)(_linalg.vector_norm)
+
 __all__ = linalg_all + _linalg.__all__
 
 del get_xp
diff --git a/array_api_compat/numpy/_aliases.py b/array_api_compat/numpy/_aliases.py
@@ -52,6 +52,7 @@
 reshape = get_xp(np)(_aliases.reshape)
 argsort = get_xp(np)(_aliases.argsort)
 sort = get_xp(np)(_aliases.sort)
+nonzero = get_xp(np)(_aliases.nonzero)
 sum = get_xp(np)(_aliases.sum)
 prod = get_xp(np)(_aliases.prod)
 ceil = get_xp(np)(_aliases.ceil)
@@ -60,8 +61,17 @@
 matmul = get_xp(np)(_aliases.matmul)
 matrix_transpose = get_xp(np)(_aliases.matrix_transpose)
 tensordot = get_xp(np)(_aliases.tensordot)
-vecdot = get_xp(np)(_aliases.vecdot)
-isdtype = get_xp(np)(_aliases.isdtype)
+
+# 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(np, 'vecdot'):
+    vecdot = np.vecdot
+else:
+    vecdot = get_xp(np)(_aliases.vecdot)
+if hasattr(np, 'isdtype'):
+    isdtype = np.isdtype
+else:
+    isdtype = get_xp(np)(_aliases.isdtype)
 
 __all__ = _aliases.__all__ + ['asarray', 'asarray_numpy', 'bool', 'acos',
                               'acosh', 'asin', 'asinh', 'atan', 'atan2',
diff --git a/array_api_compat/numpy/linalg.py b/array_api_compat/numpy/linalg.py
@@ -22,10 +22,16 @@
 pinv = get_xp(np)(_linalg.pinv)
 matrix_norm = get_xp(np)(_linalg.matrix_norm)
 svdvals = get_xp(np)(_linalg.svdvals)
-vector_norm = get_xp(np)(_linalg.vector_norm)
 diagonal = get_xp(np)(_linalg.diagonal)
 trace = get_xp(np)(_linalg.trace)
 
+# 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(np.linalg, 'vector_norm'):
+    vector_norm = np.linalg.vector_norm
+else:
+    vector_norm = get_xp(np)(_linalg.vector_norm)
+
 __all__ = linalg_all + _linalg.__all__
 
 del get_xp
diff --git a/array_api_compat/torch/_aliases.py b/array_api_compat/torch/_aliases.py
@@ -475,6 +475,8 @@ def roll(x: array, /, shift: Union[int, Tuple[int, ...]], *, axis: Optional[Unio
     return torch.roll(x, shift, axis, **kwargs)
 
 def nonzero(x: array, /, **kwargs) -> Tuple[array, ...]:
+    if x.ndim == 0:
+        raise ValueError("nonzero() does not support zero-dimensional arrays")
     return torch.nonzero(x, as_tuple=True, **kwargs)
 
 def where(condition: array, x1: array, x2: array, /) -> array:
diff --git a/array_api_compat/torch/linalg.py b/array_api_compat/torch/linalg.py
@@ -4,6 +4,8 @@
 if TYPE_CHECKING:
     import torch
     array = torch.Tensor
+    from torch import dtype as Dtype
+    from typing import Optional
 
 from torch.linalg import *
 
@@ -12,9 +14,9 @@
 from torch import linalg as torch_linalg
 linalg_all = [i for i in dir(torch_linalg) if not i.startswith('_')]
 
-# These are implemented in torch but aren't in the linalg namespace
-from torch import outer, trace
-from ._aliases import _fix_promotion, matrix_transpose, tensordot
+# outer is implemented in torch but aren't in the linalg namespace
+from torch import outer
+from ._aliases import _fix_promotion, matrix_transpose, tensordot, sum
 
 # Note: torch.linalg.cross does not default to axis=-1 (it defaults to the
 # first axis with size 3), see https://github.com/pytorch/pytorch/issues/58743
@@ -49,6 +51,11 @@ def solve(x1: array, x2: array, /, **kwargs) -> array:
     x1, x2 = _fix_promotion(x1, x2, only_scalar=False)
     return torch.linalg.solve(x1, x2, **kwargs)
 
+# torch.trace doesn't support the offset argument and doesn't support stacking
+def trace(x: array, /, *, offset: int = 0, dtype: Optional[Dtype] = None) -> array:
+    # Use our wrapped sum to make sure it does upcasting correctly
+    return sum(torch.diagonal(x, offset=offset, dim1=-2, dim2=-1), axis=-1, dtype=dtype)
+
 __all__ = linalg_all + ['outer', 'trace', 'matrix_transpose', 'tensordot',
                         'vecdot', 'solve']
 
diff --git a/cupy-xfails.txt b/cupy-xfails.txt
diff --git a/numpy-1-21-xfails.txt b/numpy-1-21-xfails.txt
diff --git a/numpy-dev-xfails.txt b/numpy-dev-xfails.txt
diff --git a/numpy-xfails.txt b/numpy-xfails.txt
diff --git a/test_cupy.sh b/test_cupy.sh
diff --git a/torch-xfails.txt b/torch-xfails.txt
