Update the function stubs from the latest version of the spec

TODO: I also need to make corresponding update to the NumPy implementation.

Author: asmeurer

array_api_tests/function_stubs/__init__.py

@@ -9,9 +9,9 @@
 
 __all__ = []
 
-from .array_object import __abs__, __add__, __and__, __array_namespace__, __bool__, __dlpack__, __dlpack_device__, __eq__, __float__, __floordiv__, __ge__, __getitem__, __gt__, __int__, __invert__, __le__, __len__, __lshift__, __lt__, __matmul__, __mod__, __mul__, __ne__, __neg__, __or__, __pos__, __pow__, __rshift__, __setitem__, __sub__, __truediv__, __xor__, __iadd__, __radd__, __iand__, __rand__, __ifloordiv__, __rfloordiv__, __ilshift__, __rlshift__, __imatmul__, __rmatmul__, __imod__, __rmod__, __imul__, __rmul__, __ior__, __ror__, __ipow__, __rpow__, __irshift__, __rrshift__, __isub__, __rsub__, __itruediv__, __rtruediv__, __ixor__, __rxor__, dtype, device, ndim, shape, size, T
+from .array_object import __abs__, __add__, __and__, __array_namespace__, __bool__, __dlpack__, __dlpack_device__, __eq__, __float__, __floordiv__, __ge__, __getitem__, __gt__, __index__, __int__, __invert__, __le__, __len__, __lshift__, __lt__, __matmul__, __mod__, __mul__, __ne__, __neg__, __or__, __pos__, __pow__, __rshift__, __setitem__, __sub__, __truediv__, __xor__, to_device, __iadd__, __radd__, __iand__, __rand__, __ifloordiv__, __rfloordiv__, __ilshift__, __rlshift__, __imatmul__, __rmatmul__, __imod__, __rmod__, __imul__, __rmul__, __ior__, __ror__, __ipow__, __rpow__, __irshift__, __rrshift__, __isub__, __rsub__, __itruediv__, __rtruediv__, __ixor__, __rxor__, dtype, device, mT, ndim, shape, size, T
 
-__all__ += ['__abs__', '__add__', '__and__', '__array_namespace__', '__bool__', '__dlpack__', '__dlpack_device__', '__eq__', '__float__', '__floordiv__', '__ge__', '__getitem__', '__gt__', '__int__', '__invert__', '__le__', '__len__', '__lshift__', '__lt__', '__matmul__', '__mod__', '__mul__', '__ne__', '__neg__', '__or__', '__pos__', '__pow__', '__rshift__', '__setitem__', '__sub__', '__truediv__', '__xor__', '__iadd__', '__radd__', '__iand__', '__rand__', '__ifloordiv__', '__rfloordiv__', '__ilshift__', '__rlshift__', '__imatmul__', '__rmatmul__', '__imod__', '__rmod__', '__imul__', '__rmul__', '__ior__', '__ror__', '__ipow__', '__rpow__', '__irshift__', '__rrshift__', '__isub__', '__rsub__', '__itruediv__', '__rtruediv__', '__ixor__', '__rxor__', 'dtype', 'device', 'ndim', 'shape', 'size', 'T']
+__all__ += ['__abs__', '__add__', '__and__', '__array_namespace__', '__bool__', '__dlpack__', '__dlpack_device__', '__eq__', '__float__', '__floordiv__', '__ge__', '__getitem__', '__gt__', '__index__', '__int__', '__invert__', '__le__', '__len__', '__lshift__', '__lt__', '__matmul__', '__mod__', '__mul__', '__ne__', '__neg__', '__or__', '__pos__', '__pow__', '__rshift__', '__setitem__', '__sub__', '__truediv__', '__xor__', 'to_device', '__iadd__', '__radd__', '__iand__', '__rand__', '__ifloordiv__', '__rfloordiv__', '__ilshift__', '__rlshift__', '__imatmul__', '__rmatmul__', '__imod__', '__rmod__', '__imul__', '__rmul__', '__ior__', '__ror__', '__ipow__', '__rpow__', '__irshift__', '__rrshift__', '__isub__', '__rsub__', '__itruediv__', '__rtruediv__', '__ixor__', '__rxor__', 'dtype', 'device', 'mT', 'ndim', 'shape', 'size', 'T']
 
 from .constants import e, inf, nan, pi
 
@@ -29,9 +29,9 @@
 
 __all__ += ['abs', 'acos', 'acosh', 'add', 'asin', 'asinh', 'atan', 'atan2', 'atanh', 'bitwise_and', 'bitwise_left_shift', 'bitwise_invert', 'bitwise_or', 'bitwise_right_shift', 'bitwise_xor', 'ceil', 'cos', 'cosh', 'divide', 'equal', 'exp', 'expm1', 'floor', 'floor_divide', 'greater', 'greater_equal', 'isfinite', 'isinf', 'isnan', 'less', 'less_equal', 'log', 'log1p', 'log2', 'log10', 'logaddexp', 'logical_and', 'logical_not', 'logical_or', 'logical_xor', 'multiply', 'negative', 'not_equal', 'positive', 'pow', 'remainder', 'round', 'sign', 'sin', 'sinh', 'square', 'sqrt', 'subtract', 'tan', 'tanh', 'trunc']
 
-from .linear_algebra_functions import einsum, matmul, tensordot, transpose, vecdot
+from .linear_algebra_functions import einsum, matmul, matrix_transpose, tensordot, vecdot
 
-__all__ += ['einsum', 'matmul', 'tensordot', 'transpose', 'vecdot']
+__all__ += ['einsum', 'matmul', 'matrix_transpose', 'tensordot', 'vecdot']
 
 from .manipulation_functions import concat, expand_dims, flip, reshape, roll, squeeze, stack
 

array_api_tests/function_stubs/array_object.py

@@ -91,6 +91,12 @@ def __gt__(self: array, other: Union[int, float, array], /) -> array:
     """
     pass
 
+def __index__(self: array, /) -> int:
+    """
+    Note: __index__ is a method of the array object.
+    """
+    pass
+
 def __int__(self: array, /) -> int:
     """
     Note: __int__ is a method of the array object.
@@ -205,6 +211,12 @@ def __xor__(self: array, other: Union[int, bool, array], /) -> array:
     """
     pass
 
+def to_device(self: array, device: device, /) -> array:
+    """
+    Note: to_device is a method of the array object.
+    """
+    pass
+
 def __iadd__(self: array, other: Union[int, float, array], /) -> array:
     """
     Note: __iadd__ is a method of the array object.
@@ -367,6 +379,9 @@ def __rxor__(self: array, other: Union[int, bool, array], /) -> array:
 # Note: device is an attribute of the array object.
 device: device = None
 
+# Note: mT is an attribute of the array object.
+mT: array = None
+
 # Note: ndim is an attribute of the array object.
 ndim: int = None
 
@@ -379,4 +394,4 @@ def __rxor__(self: array, other: Union[int, bool, array], /) -> array:
 # Note: T is an attribute of the array object.
 T: array = None
 
-__all__ = ['__abs__', '__add__', '__and__', '__array_namespace__', '__bool__', '__dlpack__', '__dlpack_device__', '__eq__', '__float__', '__floordiv__', '__ge__', '__getitem__', '__gt__', '__int__', '__invert__', '__le__', '__len__', '__lshift__', '__lt__', '__matmul__', '__mod__', '__mul__', '__ne__', '__neg__', '__or__', '__pos__', '__pow__', '__rshift__', '__setitem__', '__sub__', '__truediv__', '__xor__', '__iadd__', '__radd__', '__iand__', '__rand__', '__ifloordiv__', '__rfloordiv__', '__ilshift__', '__rlshift__', '__imatmul__', '__rmatmul__', '__imod__', '__rmod__', '__imul__', '__rmul__', '__ior__', '__ror__', '__ipow__', '__rpow__', '__irshift__', '__rrshift__', '__isub__', '__rsub__', '__itruediv__', '__rtruediv__', '__ixor__', '__rxor__', 'dtype', 'device', 'ndim', 'shape', 'size', 'T']
+__all__ = ['__abs__', '__add__', '__and__', '__array_namespace__', '__bool__', '__dlpack__', '__dlpack_device__', '__eq__', '__float__', '__floordiv__', '__ge__', '__getitem__', '__gt__', '__index__', '__int__', '__invert__', '__le__', '__len__', '__lshift__', '__lt__', '__matmul__', '__mod__', '__mul__', '__ne__', '__neg__', '__or__', '__pos__', '__pow__', '__rshift__', '__setitem__', '__sub__', '__truediv__', '__xor__', 'to_device', '__iadd__', '__radd__', '__iand__', '__rand__', '__ifloordiv__', '__rfloordiv__', '__ilshift__', '__rlshift__', '__imatmul__', '__rmatmul__', '__imod__', '__rmod__', '__imul__', '__rmul__', '__ior__', '__ror__', '__ipow__', '__rpow__', '__irshift__', '__rrshift__', '__isub__', '__rsub__', '__itruediv__', '__rtruediv__', '__ixor__', '__rxor__', 'dtype', 'device', 'mT', 'ndim', 'shape', 'size', 'T']

array_api_tests/function_stubs/linalg.py

@@ -55,6 +55,9 @@ def matrix_power(x: array, n: int, /) -> array:
 def matrix_rank(x: array, /, *, rtol: Optional[Union[float, array]] = None) -> array:
     pass
 
+def matrix_transpose(x, /):
+    pass
+
 def outer(x1: array, x2: array, /) -> array:
     pass
 
@@ -82,13 +85,10 @@ def tensordot(x1: array, x2: array, /, *, axes: Union[int, Tuple[Sequence[int],
 def trace(x: array, /, *, offset: int = 0) -> array:
     pass
 
-def transpose(x: array, /, *, axes: Optional[Tuple[int, ...]] = None) -> array:
-    pass
-
 def vecdot(x1: array, x2: array, /, *, axis: Optional[int] = None) -> array:
     pass
 
 def vector_norm(x, /, *, axis=None, keepdims=False, ord=2):
     pass
 
-__all__ = ['cholesky', 'cross', 'det', 'diagonal', 'eig', 'eigh', 'eigvals', 'eigvalsh', 'einsum', 'inv', 'matmul', 'matrix_norm', 'matrix_power', 'matrix_rank', 'outer', 'pinv', 'qr', 'slogdet', 'solve', 'svd', 'svdvals', 'tensordot', 'trace', 'transpose', 'vecdot', 'vector_norm']
+__all__ = ['cholesky', 'cross', 'det', 'diagonal', 'eig', 'eigh', 'eigvals', 'eigvalsh', 'einsum', 'inv', 'matmul', 'matrix_norm', 'matrix_power', 'matrix_rank', 'matrix_transpose', 'outer', 'pinv', 'qr', 'slogdet', 'solve', 'svd', 'svdvals', 'tensordot', 'trace', 'vecdot', 'vector_norm']

array_api_tests/function_stubs/linear_algebra_functions.py

@@ -19,13 +19,13 @@ def einsum():
 def matmul(x1: array, x2: array, /) -> array:
     pass
 
-def tensordot(x1: array, x2: array, /, *, axes: Union[int, Tuple[Sequence[int], Sequence[int]]] = 2) -> array:
+def matrix_transpose(x, /):
     pass
 
-def transpose(x: array, /, *, axes: Optional[Tuple[int, ...]] = None) -> array:
+def tensordot(x1: array, x2: array, /, *, axes: Union[int, Tuple[Sequence[int], Sequence[int]]] = 2) -> array:
     pass
 
 def vecdot(x1: array, x2: array, /, *, axis: Optional[int] = None) -> array:
     pass
 
-__all__ = ['einsum', 'matmul', 'tensordot', 'transpose', 'vecdot']
+__all__ = ['einsum', 'matmul', 'matrix_transpose', 'tensordot', 'vecdot']

array_api_tests/special_cases/test_atanh.py

@@ -7,8 +7,7 @@
 not modify it directly.
 """
 
-from ..array_helpers import (NaN, assert_exactly_equal, exactly_equal, greater, infinity, less, one,
-                             zero)
+from ..array_helpers import NaN, assert_exactly_equal, exactly_equal, greater, infinity, less, one, zero
 from ..hypothesis_helpers import numeric_arrays
 from .._array_module import atanh
 

array_api_tests/special_cases/test_sign.py

@@ -7,8 +7,7 @@
 not modify it directly.
 """
 
-from ..array_helpers import (assert_exactly_equal, exactly_equal, greater, less, logical_or, one,
-                             zero)
+from ..array_helpers import assert_exactly_equal, exactly_equal, greater, less, logical_or, one, zero
 from ..hypothesis_helpers import numeric_arrays
 from .._array_module import sign