Update stubs for the latest version of the spec

Author: asmeurer

array_api_tests/function_stubs/__init__.py

@@ -17,9 +17,9 @@
 
 __all__ += ['e', 'inf', 'nan', 'pi']
 
-from .creation_functions import arange, asarray, empty, empty_like, eye, from_dlpack, full, full_like, linspace, meshgrid, ones, ones_like, zeros, zeros_like
+from .creation_functions import arange, asarray, empty, empty_like, eye, from_dlpack, full, full_like, linspace, meshgrid, ones, ones_like, tril, triu, zeros, zeros_like
 
-__all__ += ['arange', 'asarray', 'empty', 'empty_like', 'eye', 'from_dlpack', 'full', 'full_like', 'linspace', 'meshgrid', 'ones', 'ones_like', 'zeros', 'zeros_like']
+__all__ += ['arange', 'asarray', 'empty', 'empty_like', 'eye', 'from_dlpack', 'full', 'full_like', 'linspace', 'meshgrid', 'ones', 'ones_like', 'tril', 'triu', 'zeros', 'zeros_like']
 
 from .data_type_functions import broadcast_arrays, broadcast_to, can_cast, finfo, iinfo, result_type
 
@@ -29,13 +29,13 @@
 
 __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, matrix_transpose, tensordot, vecdot
+from .linear_algebra_functions import matmul, matrix_transpose, tensordot, vecdot
 
-__all__ += ['einsum', 'matmul', 'matrix_transpose', 'tensordot', 'vecdot']
+__all__ += ['matmul', 'matrix_transpose', 'tensordot', 'vecdot']
 
-from .manipulation_functions import concat, expand_dims, flip, reshape, roll, squeeze, stack
+from .manipulation_functions import concat, expand_dims, flip, permute_dims, reshape, roll, squeeze, stack
 
-__all__ += ['concat', 'expand_dims', 'flip', 'reshape', 'roll', 'squeeze', 'stack']
+__all__ += ['concat', 'expand_dims', 'flip', 'permute_dims', 'reshape', 'roll', 'squeeze', 'stack']
 
 from .searching_functions import argmax, argmin, nonzero, where
 

array_api_tests/function_stubs/creation_functions.py

@@ -50,10 +50,16 @@ def ones(shape: Union[int, Tuple[int, ...]], *, dtype: Optional[dtype] = None, d
 def ones_like(x: array, /, *, dtype: Optional[dtype] = None, device: Optional[device] = None) -> array:
     pass
 
+def tril(x: array, /, *, k: int = 0) -> array:
+    pass
+
+def triu(x: array, /, *, k: int = 0) -> array:
+    pass
+
 def zeros(shape: Union[int, Tuple[int, ...]], *, dtype: Optional[dtype] = None, device: Optional[device] = None) -> array:
     pass
 
 def zeros_like(x: array, /, *, dtype: Optional[dtype] = None, device: Optional[device] = None) -> array:
     pass
 
-__all__ = ['arange', 'asarray', 'empty', 'empty_like', 'eye', 'from_dlpack', 'full', 'full_like', 'linspace', 'meshgrid', 'ones', 'ones_like', 'zeros', 'zeros_like']
+__all__ = ['arange', 'asarray', 'empty', 'empty_like', 'eye', 'from_dlpack', 'full', 'full_like', 'linspace', 'meshgrid', 'ones', 'ones_like', 'tril', 'triu', 'zeros', 'zeros_like']

array_api_tests/function_stubs/linalg.py

@@ -29,16 +29,13 @@ def diagonal(x: array, /, *, offset: int = 0) -> array:
 def eig():
     pass
 
-def eigh(x: array, /, *, upper: bool = False) -> Tuple[array]:
+def eigh(x: array, /) -> Tuple[array]:
     pass
 
 def eigvals():
     pass
 
-def eigvalsh(x: array, /, *, upper: bool = False) -> array:
-    pass
-
-def einsum():
+def eigvalsh(x: array, /) -> array:
     pass
 
 def inv(x: array, /) -> array:
@@ -92,4 +89,4 @@ def vecdot(x1: array, x2: array, /, *, axis: Optional[int] = None) -> array:
 def vector_norm(x: array, /, *, axis: Optional[Union[int, Tuple[int, int]]] = None, keepdims: bool = False, ord: Optional[Union[int, float, Literal[inf, -inf]]] = 2) -> array:
     pass
 
-__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']
+__all__ = ['cholesky', 'cross', 'det', 'diagonal', 'eig', 'eigh', 'eigvals', 'eigvalsh', '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

@@ -13,9 +13,6 @@
 from ._types import Optional, Tuple, Union, array
 from collections.abc import Sequence
 
-def einsum():
-    pass
-
 def matmul(x1: array, x2: array, /) -> array:
     pass
 
@@ -28,4 +25,4 @@ def tensordot(x1: array, x2: array, /, *, axes: Union[int, Tuple[Sequence[int],
 def vecdot(x1: array, x2: array, /, *, axis: Optional[int] = None) -> array:
     pass
 
-__all__ = ['einsum', 'matmul', 'matrix_transpose', 'tensordot', 'vecdot']
+__all__ = ['matmul', 'matrix_transpose', 'tensordot', 'vecdot']

array_api_tests/function_stubs/manipulation_functions.py

@@ -21,6 +21,9 @@ def expand_dims(x: array, /, *, axis: int) -> array:
 def flip(x: array, /, *, axis: Optional[Union[int, Tuple[int, ...]]] = None) -> array:
     pass
 
+def permute_dims(x: array, /, axes: Tuple[int, ...]) -> array:
+    pass
+
 def reshape(x: array, /, shape: Tuple[int, ...]) -> array:
     pass
 
@@ -33,4 +36,4 @@ def squeeze(x: array, /, axis: Union[int, Tuple[int, ...]]) -> array:
 def stack(arrays: Union[Tuple[array, ...], List[array]], /, *, axis: int = 0) -> array:
     pass
 
-__all__ = ['concat', 'expand_dims', 'flip', 'reshape', 'roll', 'squeeze', 'stack']
+__all__ = ['concat', 'expand_dims', 'flip', 'permute_dims', 'reshape', 'roll', 'squeeze', 'stack']

array_api_tests/function_stubs/statistical_functions.py

@@ -10,7 +10,7 @@
 
 from __future__ import annotations
 
-from ._types import Optional, Tuple, Union, array
+from ._types import Optional, Tuple, Union, array, dtype
 
 def max(x: array, /, *, axis: Optional[Union[int, Tuple[int, ...]]] = None, keepdims: bool = False) -> array:
     pass
@@ -21,13 +21,13 @@ def mean(x: array, /, *, axis: Optional[Union[int, Tuple[int, ...]]] = None, kee
 def min(x: array, /, *, axis: Optional[Union[int, Tuple[int, ...]]] = None, keepdims: bool = False) -> array:
     pass
 
-def prod(x: array, /, *, axis: Optional[Union[int, Tuple[int, ...]]] = None, keepdims: bool = False) -> array:
+def prod(x: array, /, *, axis: Optional[Union[int, Tuple[int, ...]]] = None, dtype: Optional[dtype] = None, keepdims: bool = False) -> array:
     pass
 
 def std(x: array, /, *, axis: Optional[Union[int, Tuple[int, ...]]] = None, correction: Union[int, float] = 0.0, keepdims: bool = False) -> array:
     pass
 
-def sum(x: array, /, *, axis: Optional[Union[int, Tuple[int, ...]]] = None, keepdims: bool = False) -> array:
+def sum(x: array, /, *, axis: Optional[Union[int, Tuple[int, ...]]] = None, dtype: Optional[dtype] = None, keepdims: bool = False) -> array:
     pass
 
 def var(x: array, /, *, axis: Optional[Union[int, Tuple[int, ...]]] = None, correction: Union[int, float] = 0.0, keepdims: bool = False) -> array: