Merge branch 'main' into more-2023

Author: asmeurer

.github/workflows/tests.yml

@@ -6,7 +6,7 @@ jobs:
     strategy:
       matrix:
         python-version: ['3.9', '3.10', '3.11', '3.12']
-        numpy-version: ['1.21', '1.26', 'dev']
+        numpy-version: ['1.21', '1.26', '2.0', 'dev']
         exclude:
           - python-version: '3.11'
             numpy-version: '1.21'

array_api_compat/common/_helpers.py

@@ -145,7 +145,7 @@ def is_ndonnx_array(x):
 
     import ndonnx as ndx
 
-    return isinstance(x, ndx.Array)  
+    return isinstance(x, ndx.Array)
 
 def is_dask_array(x):
     """
@@ -340,12 +340,9 @@ def your_function(x, y):
             elif use_compat is False:
                 namespaces.add(np)
             else:
-                # numpy 2.0 has __array_namespace__ and is fully array API
+                # numpy 2.0+ have __array_namespace__, however, they are not yet fully array API
                 # compatible.
-                if hasattr(np.empty(0), '__array_namespace__'):
-                    namespaces.add(np.empty(0).__array_namespace__(api_version=api_version))
-                else:
-                    namespaces.add(numpy_namespace)
+                namespaces.add(numpy_namespace)
         elif is_cupy_array(x):
             if _use_compat:
                 _check_api_version(api_version)
@@ -377,9 +374,13 @@ def your_function(x, y):
             elif use_compat is False:
                 import jax.numpy as jnp
             else:
-                # jax.experimental.array_api is already an array namespace. We do
-                # not have a wrapper submodule for it.
-                import jax.experimental.array_api as jnp
+                # JAX v0.4.32 and newer implements the array API directly in jax.numpy.
+                # For older JAX versions, it is available via jax.experimental.array_api.
+                import jax.numpy
+                if hasattr(jax.numpy, "__array_api_version__"):
+                    jnp = jax.numpy
+                else:
+                    import jax.experimental.array_api as jnp
             namespaces.add(jnp)
         elif is_pydata_sparse_array(x):
             if use_compat is True:
@@ -613,8 +614,9 @@ def to_device(x: Array, device: Device, /, *, stream: Optional[Union[int, Any]]
             return x
         raise ValueError(f"Unsupported device {device!r}")
     elif is_jax_array(x):
-        # This import adds to_device to x
-        import jax.experimental.array_api # noqa: F401
+        if not hasattr(x, "__array_namespace__"):
+            # In JAX v0.4.31 and older, this import adds to_device method 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

array_api_compat/torch/_aliases.py

@@ -146,6 +146,9 @@ def can_cast(from_: Union[Dtype, array], to: Dtype, /) -> bool:
 # Basic renames
 bitwise_invert = torch.bitwise_not
 newaxis = None
+# torch.conj sets the conjugation bit, which breaks conversion to other
+# libraries. See https://github.com/data-apis/array-api-compat/issues/173
+conj = torch.conj_physical
 
 # Two-arg elementwise functions
 # These require a wrapper to do the correct type promotion on 0-D tensors
@@ -707,7 +710,7 @@ def take(x: array, indices: array, /, *, axis: Optional[int] = None, **kwargs) -
     return torch.index_select(x, axis, indices, **kwargs)
 
 __all__ = ['result_type', 'can_cast', 'permute_dims', 'bitwise_invert',
-           'newaxis', 'add', 'atan2', 'bitwise_and', 'bitwise_left_shift',
+           'newaxis', 'conj', 'add', 'atan2', 'bitwise_and', 'bitwise_left_shift',
            'bitwise_or', 'bitwise_right_shift', 'bitwise_xor', 'copysign',
            'divide', 'equal', 'floor_divide', 'greater', 'greater_equal',
            'hypot', 'less', 'less_equal', 'logaddexp', 'multiply', 'not_equal',

docs/index.md

@@ -63,7 +63,8 @@ import array_api_compat.dask as da
 ```{note}
 There are no `array_api_compat` submodules for JAX, sparse, or ndonnx. These
 support for these libraries is contained in the libraries themselves (JAX
-support is in the `jax.experimental.array_api` module). The
+support is in the `jax.numpy` module in JAX v0.4.32 or newer, and in the
+`jax.experimental.array_api` module for older JAX versions). The
 array-api-compat support for these libraries consists of supporting them in
 the [helper functions](helper-functions).
 ```

numpy-dev-xfails.txt

@@ -1,32 +1,11 @@
 # finfo(float32).eps returns float32 but should return float
 array_api_tests/test_data_type_functions.py::test_finfo[float32]
 
-# NumPy deviates in some special cases for floordiv
-array_api_tests/test_special_cases.py::test_binary[floor_divide(x1_i is +infinity and isfinite(x2_i) and x2_i > 0) -> +infinity]
-array_api_tests/test_special_cases.py::test_binary[floor_divide(x1_i is +infinity and isfinite(x2_i) and x2_i < 0) -> -infinity]
-array_api_tests/test_special_cases.py::test_binary[floor_divide(x1_i is -infinity and isfinite(x2_i) and x2_i > 0) -> -infinity]
-array_api_tests/test_special_cases.py::test_binary[floor_divide(x1_i is -infinity and isfinite(x2_i) and x2_i < 0) -> +infinity]
-array_api_tests/test_special_cases.py::test_binary[floor_divide(isfinite(x1_i) and x1_i > 0 and x2_i is -infinity) -> -0]
-array_api_tests/test_special_cases.py::test_binary[floor_divide(isfinite(x1_i) and x1_i < 0 and x2_i is +infinity) -> -0]
-array_api_tests/test_special_cases.py::test_binary[__floordiv__(x1_i is +infinity and isfinite(x2_i) and x2_i > 0) -> +infinity]
-array_api_tests/test_special_cases.py::test_binary[__floordiv__(x1_i is +infinity and isfinite(x2_i) and x2_i < 0) -> -infinity]
-array_api_tests/test_special_cases.py::test_binary[__floordiv__(x1_i is -infinity and isfinite(x2_i) and x2_i > 0) -> -infinity]
-array_api_tests/test_special_cases.py::test_binary[__floordiv__(x1_i is -infinity and isfinite(x2_i) and x2_i < 0) -> +infinity]
-array_api_tests/test_special_cases.py::test_binary[__floordiv__(isfinite(x1_i) and x1_i > 0 and x2_i is -infinity) -> -0]
-array_api_tests/test_special_cases.py::test_binary[__floordiv__(isfinite(x1_i) and x1_i < 0 and x2_i is +infinity) -> -0]
-array_api_tests/test_special_cases.py::test_iop[__ifloordiv__(x1_i is +infinity and isfinite(x2_i) and x2_i > 0) -> +infinity]
-array_api_tests/test_special_cases.py::test_iop[__ifloordiv__(x1_i is +infinity and isfinite(x2_i) and x2_i < 0) -> -infinity]
-array_api_tests/test_special_cases.py::test_iop[__ifloordiv__(x1_i is -infinity and isfinite(x2_i) and x2_i > 0) -> -infinity]
-array_api_tests/test_special_cases.py::test_iop[__ifloordiv__(x1_i is -infinity and isfinite(x2_i) and x2_i < 0) -> +infinity]
-array_api_tests/test_special_cases.py::test_iop[__ifloordiv__(isfinite(x1_i) and x1_i > 0 and x2_i is -infinity) -> -0]
-array_api_tests/test_special_cases.py::test_iop[__ifloordiv__(isfinite(x1_i) and x1_i < 0 and x2_i is +infinity) -> -0]
-
 # https://github.com/numpy/numpy/issues/21213
 array_api_tests/test_special_cases.py::test_binary[__pow__(x1_i is -infinity and x2_i > 0 and not (x2_i.is_integer() and x2_i % 2 == 1)) -> +infinity]
 array_api_tests/test_special_cases.py::test_binary[__pow__(x1_i is -0 and x2_i > 0 and not (x2_i.is_integer() and x2_i % 2 == 1)) -> +0]
 array_api_tests/test_special_cases.py::test_iop[__ipow__(x1_i is -infinity and x2_i > 0 and not (x2_i.is_integer() and x2_i % 2 == 1)) -> +infinity]
 array_api_tests/test_special_cases.py::test_iop[__ipow__(x1_i is -0 and x2_i > 0 and not (x2_i.is_integer() and x2_i % 2 == 1)) -> +0]
-array_api_tests/meta/test_hypothesis_helpers.py::test_symmetric_matrices
 
 # The test suite is incorrectly checking sums that have loss of significance
 # (https://github.com/data-apis/array-api-tests/issues/168)

tests/_helpers.py

@@ -3,11 +3,8 @@
 
 import pytest
 
-wrapped_libraries = ["cupy", "torch", "dask.array"]
-all_libraries = wrapped_libraries + ["numpy", "jax.numpy", "sparse"]
-import numpy as np
-if np.__version__[0] == '1':
-    wrapped_libraries.append("numpy")
+wrapped_libraries = ["numpy", "cupy", "torch", "dask.array"]
+all_libraries = wrapped_libraries + ["jax.numpy"]
 
 # `sparse` added array API support as of Python 3.10.
 if sys.version_info >= (3, 10):
@@ -18,7 +15,11 @@ def import_(library, wrapper=False):
         pytest.importorskip(library)
     if wrapper:
         if 'jax' in library:
-            library = 'jax.experimental.array_api'
+            # JAX v0.4.32 implements the array API directly in jax.numpy
+            # Older jax versions use jax.experimental.array_api
+            jax_numpy = import_module("jax.numpy")
+            if not hasattr(jax_numpy, "__array_api_version__"):
+                library = 'jax.experimental.array_api'
         elif library.startswith('sparse'):
             library = 'sparse'
         else:

tests/test_array_namespace.py

@@ -26,8 +26,14 @@ def test_array_namespace(library, api_version, use_compat):
 
     if use_compat is False or use_compat is None and library not in wrapped_libraries:
         if library == "jax.numpy" and use_compat is None:
-            import jax.experimental.array_api
-            assert namespace == jax.experimental.array_api
+            import jax.numpy
+            if hasattr(jax.numpy, "__array_api_version__"):
+                # JAX v0.4.32 or later uses jax.numpy directly
+                assert namespace == jax.numpy
+            else:
+                # JAX v0.4.31 or earlier uses jax.experimental.array_api
+                import jax.experimental.array_api
+                assert namespace == jax.experimental.array_api
         else:
             assert namespace == xp
     else:
@@ -58,8 +64,11 @@ def test_array_namespace(library, api_version, use_compat):
 assert 'jax.experimental.array_api' not in sys.modules
 namespace = array_api_compat.array_namespace(array, api_version={api_version!r})
 
-import jax.experimental.array_api
-assert namespace == jax.experimental.array_api
+if hasattr(jax.numpy, '__array_api_version__'):
+    assert namespace == jax.numpy
+else:
+    import jax.experimental.array_api
+    assert namespace == jax.experimental.array_api
 """
         subprocess.run([sys.executable, "-c", code], check=True)