array-api-compat 1.11__tar.gz → 1.11.2__tar.gz

{array_api_compat-1.11 → array_api_compat-1.11.2}/PKG-INFO

@@ -1,6 +1,6 @@
-Metadata-Version: 2.2
+Metadata-Version: 2.4
 Name: array_api_compat
-Version: 1.11
+Version: 1.11.2
 Summary: A wrapper around NumPy and other array libraries to make them compatible with the Array API standard
 Home-page: https://data-apis.org/array-api-compat/
 Author: Consortium for Python Data API Standards
@@ -34,6 +34,7 @@ Dynamic: description
 Dynamic: description-content-type
 Dynamic: home-page
 Dynamic: license
+Dynamic: license-file
 Dynamic: provides-extra
 Dynamic: requires-python
 Dynamic: summary

{array_api_compat-1.11 → array_api_compat-1.11.2}/array_api_compat/__init__.py

@@ -17,6 +17,6 @@ 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.11'
+__version__ = '1.11.2'
 
 from .common import *  # noqa: F401, F403

{array_api_compat-1.11 → array_api_compat-1.11.2}/array_api_compat/common/_aliases.py

@@ -12,7 +12,7 @@ if TYPE_CHECKING:
 from typing import NamedTuple
 import inspect
 
-from ._helpers import array_namespace, _check_device, device, is_torch_array, is_cupy_namespace
+from ._helpers import array_namespace, _check_device, device, is_cupy_namespace
 
 # These functions are modified from the NumPy versions.
 
@@ -363,28 +363,29 @@ def clip(
 
     # At least handle the case of Python integers correctly (see
     # https://github.com/numpy/numpy/pull/26892).
-    if type(min) is int and min <= wrapped_xp.iinfo(x.dtype).min:
-        min = None
-    if type(max) is int and max >= wrapped_xp.iinfo(x.dtype).max:
-        max = None
+    if wrapped_xp.isdtype(x.dtype, "integral"):
+        if type(min) is int and min <= wrapped_xp.iinfo(x.dtype).min:
+            min = None
+        if type(max) is int and max >= wrapped_xp.iinfo(x.dtype).max:
+            max = None
 
+    dev = device(x)
     if out is None:
-        out = wrapped_xp.asarray(xp.broadcast_to(x, result_shape),
-                                 copy=True, device=device(x))
+        out = wrapped_xp.empty(result_shape, dtype=x.dtype, device=dev)
+    out[()] = x
+
     if min is not None:
-        if is_torch_array(x) and x.dtype == xp.float64 and _isscalar(min):
-            # Avoid loss of precision due to torch defaulting to float32
-            min = wrapped_xp.asarray(min, dtype=xp.float64)
-        a = xp.broadcast_to(wrapped_xp.asarray(min, device=device(x)), result_shape)
+        a = wrapped_xp.asarray(min, dtype=x.dtype, device=dev)
+        a = xp.broadcast_to(a, result_shape)
         ia = (out < a) | xp.isnan(a)
-        # torch requires an explicit cast here
-        out[ia] = wrapped_xp.astype(a[ia], out.dtype)
+        out[ia] = a[ia]
+
     if max is not None:
-        if is_torch_array(x) and x.dtype == xp.float64 and _isscalar(max):
-            max = wrapped_xp.asarray(max, dtype=xp.float64)
-        b = xp.broadcast_to(wrapped_xp.asarray(max, device=device(x)), result_shape)
+        b = wrapped_xp.asarray(max, dtype=x.dtype, device=dev)
+        b = xp.broadcast_to(b, result_shape)
         ib = (out > b) | xp.isnan(b)
-        out[ib] = wrapped_xp.astype(b[ib], out.dtype)
+        out[ib] = b[ib]
+
     # Return a scalar for 0-D
     return out[()]
 

{array_api_compat-1.11 → array_api_compat-1.11.2}/array_api_compat/cupy/_aliases.py

@@ -125,6 +125,20 @@ def astype(
     return out.copy() if copy and out is x else out
 
 
+# cupy.count_nonzero does not have keepdims
+def count_nonzero(
+    x: ndarray,
+    axis=None,
+    keepdims=False
+) -> ndarray:
+   result = cp.count_nonzero(x, axis)
+   if keepdims:
+       if axis is None:
+            return cp.reshape(result, [1]*x.ndim)
+       return cp.expand_dims(result, axis)
+   return result
+
+
 # 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'):
@@ -146,6 +160,6 @@ __all__ = _aliases.__all__ + ['__array_namespace_info__', 'asarray', 'astype',
                               'acos', 'acosh', 'asin', 'asinh', 'atan',
                               'atan2', 'atanh', 'bitwise_left_shift',
                               'bitwise_invert', 'bitwise_right_shift',
-                              'bool', 'concat', 'pow', 'sign']
+                              'bool', 'concat', 'count_nonzero', 'pow', 'sign']
 
 _all_ignore = ['cp', 'get_xp']

{array_api_compat-1.11 → array_api_compat-1.11.2}/array_api_compat/dask/array/_aliases.py

@@ -335,6 +335,21 @@ def argsort(
     return restore(x)
 
 
+# dask.array.count_nonzero does not have keepdims
+def count_nonzero(
+    x: Array,
+    axis=None,
+    keepdims=False
+) -> Array:
+   result = da.count_nonzero(x, axis)
+   if keepdims:
+       if axis is None:
+            return da.reshape(result, [1]*x.ndim)
+       return da.expand_dims(result, axis)
+   return result
+
+
+
 __all__ = _aliases.__all__ + [
                     '__array_namespace_info__', 'asarray', 'astype', 'acos',
                     'acosh', 'asin', 'asinh', 'atan', 'atan2',
@@ -343,6 +358,6 @@ __all__ = _aliases.__all__ + [
                     'result_type', 'bool', 'float32', 'float64', 'int8', 'int16', 'int32', 'int64',
                     'uint8', 'uint16', 'uint32', 'uint64',
                     'complex64', 'complex128', 'iinfo', 'finfo',
-                    'can_cast', 'result_type']
+                    'can_cast', 'count_nonzero', 'result_type']
 
 _all_ignore = ["Callable", "array_namespace", "get_xp", "da", "np"]

{array_api_compat-1.11 → array_api_compat-1.11.2}/array_api_compat/numpy/_aliases.py

@@ -127,6 +127,19 @@ def astype(
     return x.astype(dtype=dtype, copy=copy)
 
 
+# count_nonzero returns a python int for axis=None and keepdims=False
+# https://github.com/numpy/numpy/issues/17562
+def count_nonzero(
+    x : ndarray,
+    axis=None,
+    keepdims=False
+) -> ndarray:
+    result = np.count_nonzero(x, axis=axis, keepdims=keepdims)
+    if axis is None and not keepdims:
+        return np.asarray(result)
+    return result
+
+
 # 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'):
@@ -148,6 +161,6 @@ __all__ = _aliases.__all__ + ['__array_namespace_info__', 'asarray', 'astype',
                               'acos', 'acosh', 'asin', 'asinh', 'atan',
                               'atan2', 'atanh', 'bitwise_left_shift',
                               'bitwise_invert', 'bitwise_right_shift',
-                              'bool', 'concat', 'pow']
+                              'bool', 'concat', 'count_nonzero', 'pow']
 
 _all_ignore = ['np', 'get_xp']

{array_api_compat-1.11 → array_api_compat-1.11.2}/array_api_compat/torch/_aliases.py

@@ -1,15 +1,9 @@
 from __future__ import annotations
 
-from functools import wraps as _wraps
+from functools import reduce as _reduce, wraps as _wraps
 from builtins import all as _builtin_all, any as _builtin_any
 
-from ..common._aliases import (matrix_transpose as _aliases_matrix_transpose,
-                               vecdot as _aliases_vecdot,
-                               clip as _aliases_clip,
-                               unstack as _aliases_unstack,
-                               cumulative_sum as _aliases_cumulative_sum,
-                               cumulative_prod as _aliases_cumulative_prod,
-                               )
+from ..common import _aliases
 from .._internal import get_xp
 
 from ._info import __array_namespace_info__
@@ -130,25 +124,43 @@ _py_scalars = (bool, int, float, complex)
 
 
 def result_type(*arrays_and_dtypes: Union[array, Dtype, bool, int, float, complex]) -> Dtype:
-    if len(arrays_and_dtypes) == 0:
-        raise TypeError("At least one array or dtype must be provided")
-    if len(arrays_and_dtypes) == 1:
+    num = len(arrays_and_dtypes)
+
+    if num == 0:
+        raise ValueError("At least one array or dtype must be provided")
+
+    elif num == 1:
         x = arrays_and_dtypes[0]
         if isinstance(x, torch.dtype):
             return x
         return x.dtype
-    if len(arrays_and_dtypes) > 2:
-        return result_type(arrays_and_dtypes[0], result_type(*arrays_and_dtypes[1:]))
 
-    x, y = arrays_and_dtypes
-    if isinstance(x, _py_scalars) or isinstance(y, _py_scalars):
-        return torch.result_type(x, y)
+    if num == 2:
+        x, y = arrays_and_dtypes
+        return _result_type(x, y)
+
+    else:
+        # sort scalars so that they are treated last
+        scalars, others = [], []
+        for x in arrays_and_dtypes:
+            if isinstance(x, _py_scalars):
+                scalars.append(x)
+            else:
+                others.append(x)
+        if not others:
+            raise ValueError("At least one array or dtype must be provided")
+
+        # combine left-to-right
+        return _reduce(_result_type, others + scalars)
 
-    xdt = x.dtype if not isinstance(x, torch.dtype) else x
-    ydt = y.dtype if not isinstance(y, torch.dtype) else y
 
-    if (xdt, ydt) in _promotion_table:
-        return _promotion_table[xdt, ydt]
+def _result_type(x, y):
+    if not (isinstance(x, _py_scalars) or isinstance(y, _py_scalars)):
+        xdt = x.dtype if not isinstance(x, torch.dtype) else x
+        ydt = y.dtype if not isinstance(y, torch.dtype) else y
+
+        if (xdt, ydt) in _promotion_table:
+            return _promotion_table[xdt, ydt]
 
     # This doesn't result_type(dtype, dtype) for non-array API dtypes
     # because torch.result_type only accepts tensors. This does however, allow
@@ -157,6 +169,7 @@ def result_type(*arrays_and_dtypes: Union[array, Dtype, bool, int, float, comple
     y = torch.tensor([], dtype=y) if isinstance(y, torch.dtype) else y
     return torch.result_type(x, y)
 
+
 def can_cast(from_: Union[Dtype, array], to: Dtype, /) -> bool:
     if not isinstance(from_, torch.dtype):
         from_ = from_.dtype
@@ -215,10 +228,10 @@ def min(x: array, /, *, axis: Optional[Union[int, Tuple[int, ...]]] = None, keep
         return torch.clone(x)
     return torch.amin(x, axis, keepdims=keepdims)
 
-clip = get_xp(torch)(_aliases_clip)
-unstack = get_xp(torch)(_aliases_unstack)
-cumulative_sum = get_xp(torch)(_aliases_cumulative_sum)
-cumulative_prod = get_xp(torch)(_aliases_cumulative_prod)
+clip = get_xp(torch)(_aliases.clip)
+unstack = get_xp(torch)(_aliases.unstack)
+cumulative_sum = get_xp(torch)(_aliases.cumulative_sum)
+cumulative_prod = get_xp(torch)(_aliases.cumulative_prod)
 
 # torch.sort also returns a tuple
 # https://github.com/pytorch/pytorch/issues/70921
@@ -527,7 +540,7 @@ def diff(
     return torch.diff(x, dim=axis, n=n, prepend=prepend, append=append)
 
 
-# torch uses `dim` instead of `axis`
+# torch uses `dim` instead of `axis`, does not have keepdims
 def count_nonzero(
     x: array,
     /,
@@ -535,7 +548,14 @@ def count_nonzero(
     axis: Optional[Union[int, Tuple[int, ...]]] = None,
     keepdims: bool = False,
 ) -> array:
-    return torch.count_nonzero(x, dim=axis, keepdims=keepdims)
+    result = torch.count_nonzero(x, dim=axis)
+    if keepdims:
+        if axis is not None:
+            return result.unsqueeze(axis)
+        return _axis_none_keepdims(result, x.ndim, keepdims)
+    else:
+        return result
+
 
 
 def where(condition: array, x1: array, x2: array, /) -> array:
@@ -710,8 +730,8 @@ def matmul(x1: array, x2: array, /, **kwargs) -> array:
     x1, x2 = _fix_promotion(x1, x2, only_scalar=False)
     return torch.matmul(x1, x2, **kwargs)
 
-matrix_transpose = get_xp(torch)(_aliases_matrix_transpose)
-_vecdot = get_xp(torch)(_aliases_vecdot)
+matrix_transpose = get_xp(torch)(_aliases.matrix_transpose)
+_vecdot = get_xp(torch)(_aliases.vecdot)
 
 def vecdot(x1: array, x2: array, /, *, axis: int = -1) -> array:
     x1, x2 = _fix_promotion(x1, x2, only_scalar=False)

{array_api_compat-1.11 → array_api_compat-1.11.2}/array_api_compat.egg-info/PKG-INFO

@@ -1,6 +1,6 @@
-Metadata-Version: 2.2
+Metadata-Version: 2.4
 Name: array_api_compat
-Version: 1.11
+Version: 1.11.2
 Summary: A wrapper around NumPy and other array libraries to make them compatible with the Array API standard
 Home-page: https://data-apis.org/array-api-compat/
 Author: Consortium for Python Data API Standards
@@ -34,6 +34,7 @@ Dynamic: description
 Dynamic: description-content-type
 Dynamic: home-page
 Dynamic: license
+Dynamic: license-file
 Dynamic: provides-extra
 Dynamic: requires-python
 Dynamic: summary

{array_api_compat-1.11 → array_api_compat-1.11.2}/array_api_compat.egg-info/SOURCES.txt

@@ -44,4 +44,5 @@ tests/test_dask.py
 tests/test_isdtype.py
 tests/test_jax.py
 tests/test_no_dependencies.py
+tests/test_torch.py
 tests/test_vendoring.py

{array_api_compat-1.11 → array_api_compat-1.11.2}/tests/test_all.py

@@ -16,6 +16,7 @@ from ._helpers import import_, wrapped_libraries
 
 import pytest
 
+@pytest.mark.skip(reason="TODO: starts failing after adding test_torch.py in gh-277")
 @pytest.mark.parametrize("library", ["common"] + wrapped_libraries)
 def test_all(library):
     if library == "common":

{array_api_compat-1.11 → array_api_compat-1.11.2}/tests/test_common.py

@@ -367,3 +367,18 @@ def test_asarray_copy(library):
         assert all(b[0] == 1.0)
     else:
         assert all(b[0] == 0.0)
+
+
+@pytest.mark.parametrize("library", ["numpy", "cupy", "torch"])
+def test_clip_out(library):
+    """Test non-standard out= parameter for clip()
+
+    (see "Avoid Restricting Behavior that is Outside the Scope of the Standard"
+    in https://data-apis.org/array-api-compat/dev/special-considerations.html)
+    """
+    xp = import_(library, wrapper=True)
+    x = xp.asarray([10, 20, 30])
+    out = xp.zeros_like(x)
+    xp.clip(x, 15, 25, out=out)
+    expect = xp.asarray([15, 20, 25])
+    assert xp.all(out == expect)

array_api_compat-1.11.2/tests/test_torch.py

@@ -0,0 +1,98 @@
+"""Test "unspecified" behavior which we cannot easily test in the Array API test suite.
+"""
+import itertools
+
+import pytest
+import torch
+
+from array_api_compat import torch as xp
+
+
+class TestResultType:
+    def test_empty(self):
+        with pytest.raises(ValueError):
+            xp.result_type()
+
+    def test_one_arg(self):
+        for x in [1, 1.0, 1j, '...', None]:
+            with pytest.raises((ValueError, AttributeError)):
+                xp.result_type(x)
+
+        for x in [xp.float32, xp.int64, torch.complex64]:
+            assert xp.result_type(x) == x
+
+        for x in [xp.asarray(True, dtype=xp.bool), xp.asarray(1, dtype=xp.complex64)]:
+            assert xp.result_type(x) == x.dtype
+
+    def test_two_args(self):
+        # Only include here things "unspecified" in the spec
+
+        # scalar, tensor or tensor,tensor
+        for x, y in [
+            (1., 1j),
+            (1j, xp.arange(3)),
+            (True, xp.asarray(3.)),
+            (xp.ones(3) == 1, 1j*xp.ones(3)),
+        ]:
+            assert xp.result_type(x, y) == torch.result_type(x, y)
+
+        # dtype, scalar
+        for x, y in [
+            (1j, xp.int64),
+            (True, xp.float64),
+        ]:
+            assert xp.result_type(x, y) == torch.result_type(x, xp.empty([], dtype=y))
+
+        # dtype, dtype
+        for x, y in [
+            (xp.bool, xp.complex64)
+        ]:
+            xt, yt = xp.empty([], dtype=x), xp.empty([], dtype=y)
+            assert xp.result_type(x, y) == torch.result_type(xt, yt)
+
+    def test_multi_arg(self):
+        torch.set_default_dtype(torch.float32)
+
+        args = [1., 5, 3, torch.asarray([3], dtype=torch.float16), 5, 6, 1.]
+        assert xp.result_type(*args) == torch.float16
+
+        args = [1, 2, 3j, xp.arange(3, dtype=xp.float32), 4, 5, 6]
+        assert xp.result_type(*args) == xp.complex64
+
+        args = [1, 2, 3j, xp.float64, 4, 5, 6]
+        assert xp.result_type(*args) == xp.complex128
+
+        args = [1, 2, 3j, xp.float64, 4, xp.asarray(3, dtype=xp.int16), 5, 6, False]
+        assert xp.result_type(*args) == xp.complex128
+
+        i64 = xp.ones(1, dtype=xp.int64)
+        f16 = xp.ones(1, dtype=xp.float16)
+        for i in itertools.permutations([i64, f16, 1.0, 1.0]):
+            assert xp.result_type(*i) == xp.float16, f"{i}"
+
+        with pytest.raises(ValueError):
+            xp.result_type(1, 2, 3, 4)
+
+
+    @pytest.mark.parametrize("default_dt", ['float32', 'float64'])
+    @pytest.mark.parametrize("dtype_a",
+        (xp.int32, xp.int64, xp.float32, xp.float64, xp.complex64, xp.complex128)
+    )
+    @pytest.mark.parametrize("dtype_b", 
+        (xp.int32, xp.int64, xp.float32, xp.float64, xp.complex64, xp.complex128)
+    )
+    def test_gh_273(self, default_dt, dtype_a, dtype_b):
+        # Regression test for https://github.com/data-apis/array-api-compat/issues/273
+
+        try:
+            prev_default = torch.get_default_dtype()
+            default_dtype = getattr(torch, default_dt)
+            torch.set_default_dtype(default_dtype)
+
+            a = xp.asarray([2, 1], dtype=dtype_a)
+            b = xp.asarray([1, -1], dtype=dtype_b)
+            dtype_1 = xp.result_type(a, b, 1.0)
+            dtype_2 = xp.result_type(b, a, 1.0)
+            assert dtype_1 == dtype_2
+        finally:
+            torch.set_default_dtype(prev_default)