Merge pull request #38 from asmeurer/2022.12

2022.12 support

Author: asmeurer

CHANGELOG.md

@@ -1,3 +1,19 @@
+# 1.3 (2023-06-20)
+
+## Major Changes
+
+- Add [2022.12](https://data-apis.org/array-api/2022.12/) standard support.
+  This includes things like adding complex dtype support, adding the new
+  `take` function, and various minor changes in the specification.
+
+## Minor Changes
+
+- Support `"cpu"` in CuPy `to_device()`.
+
+- Return a new array in NumPy/CuPy `reshape(copy=False)`.
+
+- Fix signatures for PyTorch `broadcast_to` and `permute_dims`.
+
 # 1.2 (2023-04-03)
 
 ## Major Changes

README.md

@@ -12,11 +12,8 @@ each array library itself fully compatible with the array API, but this
 requires making backwards incompatible changes in many cases, so this will
 take some time.
 
-Currently all libraries here are implemented against the 2021.12 version of
-the standard. Support for the [2022.12
-version](https://data-apis.org/array-api/2022.12/changelog.html), which adds
-complex number support as well as several additional functions, will be added
-later this year.
+Currently all libraries here are implemented against the [2022.22
+version](https://data-apis.org/array-api/2022.22/) of the standard.
 
 ## Usage
 
@@ -177,8 +174,6 @@ version.
   in the spec. Use the `size(x)` helper function as a portable workaround (see
   above).
 
-- The `linalg` extension is not yet implemented.
-
 - PyTorch does not have unsigned integer types other than `uint8`, and no
   attempt is made to implement them here.
 

array_api_compat/__init__.py

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

array_api_compat/common/_aliases.py

@@ -397,9 +397,12 @@ def sum(
     keepdims: bool = False,
     **kwargs,
 ) -> ndarray:
-    # `xp.sum` already upcasts integers, but not floats
-    if dtype is None and x.dtype == xp.float32:
-        dtype = xp.float64
+    # `xp.sum` already upcasts integers, but not floats or complexes
+    if dtype is None:
+        if x.dtype == xp.float32:
+            dtype = xp.float64
+        elif x.dtype == xp.complex64:
+            dtype = xp.complex128
     return xp.sum(x, axis=axis, dtype=dtype, keepdims=keepdims, **kwargs)
 
 def prod(
@@ -412,8 +415,11 @@ def prod(
     keepdims: bool = False,
     **kwargs,
 ) -> ndarray:
-    if dtype is None and x.dtype == xp.float32:
-        dtype = xp.float64
+    if dtype is None:
+        if x.dtype == xp.float32:
+            dtype = xp.float64
+        elif x.dtype == xp.complex64:
+            dtype = xp.complex128
     return xp.prod(x, dtype=dtype, axis=axis, keepdims=keepdims, **kwargs)
 
 # ceil, floor, and trunc return integers for integer inputs

array_api_compat/common/_linalg.py

@@ -136,8 +136,13 @@ def vector_norm(x: ndarray, /, xp, *, axis: Optional[Union[int, Tuple[int, ...]]
 def diagonal(x: ndarray, /, xp, *, offset: int = 0, **kwargs) -> ndarray:
     return xp.diagonal(x, offset=offset, axis1=-2, axis2=-1, **kwargs)
 
-def trace(x: ndarray, /, xp, *, offset: int = 0, **kwargs) -> ndarray:
-    return xp.asarray(xp.trace(x, offset=offset, axis1=-2, axis2=-1, **kwargs))
+def trace(x: ndarray, /, xp, *, offset: int = 0, dtype=None, **kwargs) -> ndarray:
+    if dtype is None:
+        if x.dtype == xp.float32:
+            dtype = xp.float64
+        elif x.dtype == xp.complex64:
+            dtype = xp.complex128
+    return xp.asarray(xp.trace(x, offset=offset, dtype=dtype, axis1=-2, axis2=-1, **kwargs))
 
 __all__ = ['cross', 'matmul', 'outer', 'tensordot', 'EighResult',
            'QRResult', 'SlogdetResult', 'SVDResult', 'eigh', 'qr', 'slogdet',

array_api_compat/cupy/__init__.py

@@ -13,4 +13,4 @@
 
 from ..common._helpers import *
 
-__array_api_version__ = '2021.12'
+__array_api_version__ = '2022.12'

array_api_compat/numpy/__init__.py

@@ -19,4 +19,4 @@
 
 from ..common._helpers import *
 
-__array_api_version__ = '2021.12'
+__array_api_version__ = '2022.12'

array_api_compat/torch/__init__.py

@@ -19,4 +19,4 @@
 
 from ..common._helpers import *
 
-__array_api_version__ = '2021.12'
+__array_api_version__ = '2022.12'

array_api_compat/torch/_aliases.py

@@ -32,6 +32,8 @@
     *_int_dtypes,
     torch.float32,
     torch.float64,
+    torch.complex64,
+    torch.complex128,
 }
 
 _promotion_table  = {
@@ -70,6 +72,16 @@
     (torch.float32, torch.float64): torch.float64,
     (torch.float64, torch.float32): torch.float64,
     (torch.float64, torch.float64): torch.float64,
+    # complexes
+    (torch.complex64, torch.complex64): torch.complex64,
+    (torch.complex64, torch.complex128): torch.complex128,
+    (torch.complex128, torch.complex64): torch.complex128,
+    (torch.complex128, torch.complex128): torch.complex128,
+    # Mixed float and complex
+    (torch.float32, torch.complex64): torch.complex64,
+    (torch.float32, torch.complex128): torch.complex128,
+    (torch.float64, torch.complex64): torch.complex128,
+    (torch.float64, torch.complex128): torch.complex128,
 }
 
 
@@ -129,7 +141,6 @@ def can_cast(from_: Union[Dtype, array], to: Dtype, /) -> bool:
     return torch.can_cast(from_, to)
 
 # Basic renames
-permute_dims = torch.permute
 bitwise_invert = torch.bitwise_not
 
 # Two-arg elementwise functions
@@ -439,18 +450,26 @@ def squeeze(x: array, /, axis: Union[int, Tuple[int, ...]]) -> array:
         x = torch.squeeze(x, a)
     return x
 
+# torch.broadcast_to uses size instead of shape
+def broadcast_to(x: array, /, shape: Tuple[int, ...], **kwargs) -> array:
+    return torch.broadcast_to(x, shape, **kwargs)
+
+# torch.permute uses dims instead of axes
+def permute_dims(x: array, /, axes: Tuple[int, ...]) -> array:
+    return torch.permute(x, axes)
+
 # The axis parameter doesn't work for flip() and roll()
 # https://github.com/pytorch/pytorch/issues/71210. Also torch.flip() doesn't
 # accept axis=None
-def flip(x: array, /, *, axis: Optional[Union[int, Tuple[int, ...]]] = None) -> array:
+def flip(x: array, /, *, axis: Optional[Union[int, Tuple[int, ...]]] = None, **kwargs) -> array:
     if axis is None:
         axis = tuple(range(x.ndim))
     # torch.flip doesn't accept dim as an int but the method does
     # https://github.com/pytorch/pytorch/issues/18095
-    return x.flip(axis)
+    return x.flip(axis, **kwargs)
 
-def roll(x: array, /, shift: Union[int, Tuple[int, ...]], *, axis: Optional[Union[int, Tuple[int, ...]]] = None) -> array:
-    return torch.roll(x, shift, axis)
+def roll(x: array, /, shift: Union[int, Tuple[int, ...]], *, axis: Optional[Union[int, Tuple[int, ...]]] = None, **kwargs) -> array:
+    return torch.roll(x, shift, axis, **kwargs)
 
 def nonzero(x: array, /, **kwargs) -> Tuple[array, ...]:
     return torch.nonzero(x, as_tuple=True, **kwargs)
@@ -662,15 +681,18 @@ def isdtype(
     else:
         return dtype == kind
 
+def take(x: array, indices: array, /, *, axis: int, **kwargs) -> array:
+    return torch.index_select(x, axis, indices, **kwargs)
+
 __all__ = ['result_type', 'can_cast', 'permute_dims', 'bitwise_invert', 'add',
            'atan2', 'bitwise_and', 'bitwise_left_shift', 'bitwise_or',
            'bitwise_right_shift', 'bitwise_xor', 'divide', 'equal',
            'floor_divide', 'greater', 'greater_equal', 'less', 'less_equal',
            'logaddexp', 'multiply', 'not_equal', 'pow', 'remainder',
            'subtract', 'max', 'min', 'sort', 'prod', 'sum', 'any', 'all',
-           'mean', 'std', 'var', 'concat', 'squeeze', 'flip', 'roll',
+           'mean', 'std', 'var', 'concat', 'squeeze', 'broadcast_to', 'flip', 'roll',
            'nonzero', 'where', 'reshape', 'arange', 'eye', 'linspace', 'full',
            'ones', 'zeros', 'empty', 'tril', 'triu', 'expand_dims', 'astype',
            'broadcast_arrays', 'unique_all', 'unique_counts',
            'unique_inverse', 'unique_values', 'matmul', 'matrix_transpose',
-           'vecdot', 'tensordot', 'isdtype']
+           'vecdot', 'tensordot', 'isdtype', 'take']

cupy-xfails.txt

@@ -41,6 +41,7 @@ array_api_tests/test_operators_and_elementwise_functions.py::test_remainder[__mo
 array_api_tests/test_operators_and_elementwise_functions.py::test_remainder[__imod__(x1, x2)]
 array_api_tests/test_operators_and_elementwise_functions.py::test_remainder[__mod__(x, s)]
 array_api_tests/test_operators_and_elementwise_functions.py::test_subtract[__sub__(x, s)]
+array_api_tests/test_operators_and_elementwise_functions.py::test_add[__add__(x, s)]
 # floating point inaccuracy
 array_api_tests/test_operators_and_elementwise_functions.py::test_remainder[remainder(x1, x2)]
 
@@ -55,6 +56,10 @@ array_api_tests/test_statistical_functions.py::test_max
 # (https://github.com/data-apis/array-api-tests/issues/171)
 array_api_tests/test_signatures.py::test_func_signature[meshgrid]
 
+# testsuite issue with test_square
+# https://github.com/data-apis/array-api-tests/issues/190
+array_api_tests/test_operators_and_elementwise_functions.py::test_square
+
 # We cannot add array attributes
 array_api_tests/test_signatures.py::test_array_method_signature[__array_namespace__]
 array_api_tests/test_signatures.py::test_array_method_signature[__index__]