From bc8d9a6a86931c84d82c27639e721bb9a28a2262 Mon Sep 17 00:00:00 2001 From: Willow Ahrens Date: Thu, 29 Oct 2026 12:58:29 -0400 Subject: [PATCH 1/6] add reduce/elemwise --- src/array_api_stubs/_draft/reduce.py | 189 +++++++++++++++++++++++++++ 1 file changed, 189 insertions(+) create mode 100644 src/array_api_stubs/_draft/reduce.py diff --git a/src/array_api_stubs/_draft/reduce.py b/src/array_api_stubs/_draft/reduce.py new file mode 100644 index 000000000..1ef310535 --- /dev/null +++ b/src/array_api_stubs/_draft/reduce.py @@ -0,0 +1,189 @@ + +nary_ops = { + "+": "add", + "add": "add", + "-": "subtract", + "sub": "subtract", + "subtract": "subtract", + "*": "multiply", + "mul": "multiply", + "multiply": "multiply", + "/": "divide", + "div": "divide", + "divide": "divide", + "//": "floor_divide", + "fld": "floor_divide", + "floor_divide": "floor_divide", + "%": "remainder", + "mod": "remainder", + "remainder": "remainder", + "**": "power", + "pow": "power", + "power": "power", + "==": "equal", + "eq": "equal", + "equal": "equal", + "!=": "not_equal", + "ne": "not_equal", + "not_equal": "not_equal", + "<": "less", + "lt": "less", + "less": "less", + "<=": "less_equal", + "le": "less_equal", + "less_equal": "less_equal", + ">": "greater", + "gt": "greater", + "greater": "greater", + ">=": "greater_equal", + "ge": "greater_equal", + "greater_equal": "greater_equal", + "&": "bitwise_and", + "bitwise_and": "bitwise_and", + "|": "bitwise_or", + "bitwise_or": "bitwise_or", + "^": "bitwise_xor", + "bitwise_xor": "bitwise_xor", + "<<": "bitwise_left_shift", + "lshift": "bitwise_left_shift", + "bitwise_left_shift": "bitwise_left_shift", + ">>": "bitwise_right_shift", + "rshift": "bitwise_right_shift", + "bitwise_right_shift": "bitwise_right_shift", + "and": "logical_and", + "or": "logical_or", + "not": "logical_not", + "min": "minimum", + "max": "maximum", + "logaddexp": "logaddexp", +} + + +unary_ops = { + "+": "positive", + "pos": "positive", + "positive": "positive", + "-": "negative", + "neg": "negative", + "negative": "negative", + "~": "bitwise_invert", + "invert": "bitwise_invert", + "bitwise_invert": "bitwise_invert", + "not": "logical_not", + "logical_not": "logical_not", + "abs": "absolute", + "absolute": "absolute", + "sqrt": "sqrt", + "exp": "exp", + "log": "log", + "log1p": "log1p", + "log10": "log10", + "log2": "log2", + "sin": "sin", + "cos": "cos", + "tan": "tan", + "sinh": "sinh", + "cosh": "cosh", + "tanh": "tanh", + "asin": "arcsin", + "acos": "arccos", + "atan": "arctan", + "asinh": "arcsinh", + "acosh": "arccosh", + "atanh": "arctanh", +} + + +reduction_ops = { + "+": "sum", + "add": "sum", + "sum": "sum", + "*": "prod", + "mul": "prod", + "prod": "prod", + "and": "all", + "or": "any", + "min": "min", + "max": "max", + "argmin": "argmin", + "argmax": "argmax", + "mean": "mean", + "std": "std", + "var": "var", + "count_nonzero": "count_nonzero", + # "&": "bitwise_and", + # "|": "bitwise_or", + # "^": "bitwise_xor", +} + +def reduce( + op: FinchOperator, + x, + /, + *, + axis: int | tuple[int, ...] | None = None, + dtype=None, + keepdims: bool = False, + init=None, +): + """ + Reduces the input array ``x`` with the binary operator ``op``. Reduces along + the specified `axis`, with an initial value `init`. + + Parameters + ---------- + x: array + input array. Should have a numeric data type. + axis: Optional[Union[int, Tuple[int, ...]]] + axis or axes along which reduction must be computed. By default, the reduction + must be computed over the entire array. If a tuple of integers, reductions must + be computed over multiple axes. Default: ``None``. + + dtype: Optional[dtype] + data type of the returned array. If ``None``, a suitable data type will be + calculated. + + keepdims: bool + if ``True``, the reduced axes (dimensions) must be included in the result as + singleton dimensions, and, accordingly, the result must be compatible with the + input array (see :ref:`broadcasting`). Otherwise, if ``False``, the reduced axes + (dimensions) must not be included in the result. Default: ``False``. + + init: Optional + Initial value for the reduction. If ``None``, a suitable initial value will be + calculated. The initial value must be compatible with the operation defined by + ``op``. For example, if ``op`` is addition, the initial value should be zero; if + ``op`` is multiplication, the initial value should be one. + + Returns + ------- + out: array + If the reduction was computed over the entire array, a zero-dimensional array + containing the reduction; otherwise, a non-zero-dimensional array containing the + reduction. The returned array must have a data type as described by the + ``dtype`` parameter above. + """ + +def elementwise(f: FinchOperator, *args) -> LazyTensor: + """ + elementwise(f, *args) -> LazyTensor: + + Applies the function f elementwise to the given arguments, following + [broacasting rules](https://numpy.org/doc/stable/user/basics.broadcasting.html). + The function f should be a callable that takes the same number of arguments + as the number of tensors passed to `elementwise`. + + The function will automatically handle broadcasting of the input tensors to + ensure they have compatible shapes. For example, `elementwise(ffunc.add, + x, y)` is equivalent to `x + y`. + + Parameters: + - f: The function to apply elementwise. + - *args: The tensors to apply the function to. These tensors should be + compatible for broadcasting. + + Returns: + - LazyTensor: The tensor, `out`, of results from applying `f` elementwise to + the input tensors. After broadcasting the arguments to the same shape, for + each index `i`, `out[*i] = f(args[0][*i], args[1][*i], ...)`. + """ From 64a1e838e1dda82c958f0b7848c4da616d8bbe70 Mon Sep 17 00:00:00 2001 From: Willow Ahrens Date: Thu, 29 Oct 2026 13:16:26 -0400 Subject: [PATCH 2/6] higher-order functions --- spec/draft/API_specification/higher_order.rst | 22 +++++ spec/draft/API_specification/index.rst | 1 + src/array_api_stubs/_draft/__init__.py | 1 + src/array_api_stubs/_draft/_types.py | 2 + .../_draft/{reduce.py => higher_order.py} | 86 +++++++++++-------- 5 files changed, 74 insertions(+), 38 deletions(-) create mode 100644 spec/draft/API_specification/higher_order.rst rename src/array_api_stubs/_draft/{reduce.py => higher_order.py} (70%) diff --git a/spec/draft/API_specification/higher_order.rst b/spec/draft/API_specification/higher_order.rst new file mode 100644 index 000000000..2fb57b318 --- /dev/null +++ b/spec/draft/API_specification/higher_order.rst @@ -0,0 +1,22 @@ +Higher-order Functions +====================== + + Array API specification for higher-order functions. + +A conforming implementation of the array API standard must provide and support the following functions. + + +Objects in API +-------------- + +.. currentmodule:: array_api + +.. + NOTE: please keep the functions in alphabetical order + +.. autosummary:: + :toctree: generated + :template: method.rst + + elementwise + reduce diff --git a/spec/draft/API_specification/index.rst b/spec/draft/API_specification/index.rst index ffc3d3775..190573bdd 100644 --- a/spec/draft/API_specification/index.rst +++ b/spec/draft/API_specification/index.rst @@ -27,6 +27,7 @@ A conforming implementation of the array API standard must provide and support t data_types elementwise_functions function_and_method_signatures + higher_order indexing indexing_functions inspection diff --git a/src/array_api_stubs/_draft/__init__.py b/src/array_api_stubs/_draft/__init__.py index 537ea8f85..78a47f1ba 100644 --- a/src/array_api_stubs/_draft/__init__.py +++ b/src/array_api_stubs/_draft/__init__.py @@ -6,6 +6,7 @@ from .data_type_functions import * from . import data_types as dtype from .elementwise_functions import * +from .higher_order import * from .indexing_functions import * from .linear_algebra_functions import * from .manipulation_functions import * diff --git a/src/array_api_stubs/_draft/_types.py b/src/array_api_stubs/_draft/_types.py index f2fa356f2..0c859280a 100644 --- a/src/array_api_stubs/_draft/_types.py +++ b/src/array_api_stubs/_draft/_types.py @@ -8,6 +8,7 @@ __all__ = [ "Any", + "Callable", "List", "Literal", "NestedSequence", @@ -34,6 +35,7 @@ from dataclasses import dataclass from typing import ( Any, + Callable, List, Literal, Optional, diff --git a/src/array_api_stubs/_draft/reduce.py b/src/array_api_stubs/_draft/higher_order.py similarity index 70% rename from src/array_api_stubs/_draft/reduce.py rename to src/array_api_stubs/_draft/higher_order.py index 1ef310535..93054e9f1 100644 --- a/src/array_api_stubs/_draft/reduce.py +++ b/src/array_api_stubs/_draft/higher_order.py @@ -1,5 +1,9 @@ +__all__ = ["reduce", "elementwise"] -nary_ops = { +from ._types import Callable, Optional, Tuple, Union, array, dtype + + +binary_ops = { "+": "add", "add": "add", "-": "subtract", @@ -17,9 +21,9 @@ "%": "remainder", "mod": "remainder", "remainder": "remainder", - "**": "power", - "pow": "power", - "power": "power", + "**": "pow", + "pow": "pow", + "power": "pow", "==": "equal", "eq": "equal", "equal": "equal", @@ -52,7 +56,6 @@ "bitwise_right_shift": "bitwise_right_shift", "and": "logical_and", "or": "logical_or", - "not": "logical_not", "min": "minimum", "max": "maximum", "logaddexp": "logaddexp", @@ -71,8 +74,8 @@ "bitwise_invert": "bitwise_invert", "not": "logical_not", "logical_not": "logical_not", - "abs": "absolute", - "absolute": "absolute", + "abs": "abs", + "absolute": "abs", "sqrt": "sqrt", "exp": "exp", "log": "log", @@ -85,12 +88,12 @@ "sinh": "sinh", "cosh": "cosh", "tanh": "tanh", - "asin": "arcsin", - "acos": "arccos", - "atan": "arctan", - "asinh": "arcsinh", - "acosh": "arccosh", - "atanh": "arctanh", + "asin": "asin", + "acos": "acos", + "atan": "atan", + "asinh": "asinh", + "acosh": "acosh", + "atanh": "atanh", } @@ -116,22 +119,25 @@ # "^": "bitwise_xor", } + def reduce( - op: FinchOperator, - x, + op: Callable[[array, array], array], + x: array, /, *, - axis: int | tuple[int, ...] | None = None, - dtype=None, + axis: Optional[Union[int, Tuple[int, ...]]] = None, + dtype: Optional[dtype] = None, keepdims: bool = False, - init=None, -): + init: Optional[Union[bool, int, float, complex]] = None, +) -> array: """ Reduces the input array ``x`` with the binary operator ``op``. Reduces along - the specified `axis`, with an initial value `init`. + the specified ``axis``, with an initial value ``init``. Parameters ---------- + op: Callable[[array, array], array] + binary operator with which to perform the reduction (e.g., :func:`~array_api.add`). x: array input array. Should have a numeric data type. axis: Optional[Union[int, Tuple[int, ...]]] @@ -149,7 +155,7 @@ def reduce( input array (see :ref:`broadcasting`). Otherwise, if ``False``, the reduced axes (dimensions) must not be included in the result. Default: ``False``. - init: Optional + init: Optional[Union[bool, int, float, complex]] Initial value for the reduction. If ``None``, a suitable initial value will be calculated. The initial value must be compatible with the operation defined by ``op``. For example, if ``op`` is addition, the initial value should be zero; if @@ -164,26 +170,30 @@ def reduce( ``dtype`` parameter above. """ -def elementwise(f: FinchOperator, *args) -> LazyTensor: - """ - elementwise(f, *args) -> LazyTensor: - Applies the function f elementwise to the given arguments, following - [broacasting rules](https://numpy.org/doc/stable/user/basics.broadcasting.html). - The function f should be a callable that takes the same number of arguments - as the number of tensors passed to `elementwise`. +def elementwise(f: Callable[..., array], /, *args: array) -> array: + """ + Applies the function ``f`` elementwise to the given arguments, following + broadcasting rules (see :ref:`broadcasting`). + The function ``f`` should be a callable that takes the same number of arguments + as the number of arrays passed to ``elementwise``. - The function will automatically handle broadcasting of the input tensors to - ensure they have compatible shapes. For example, `elementwise(ffunc.add, - x, y)` is equivalent to `x + y`. + The function will automatically handle broadcasting of the input arrays to + ensure they have compatible shapes. For example, ``elementwise(add, x, y)`` + is equivalent to ``x + y``. - Parameters: - - f: The function to apply elementwise. - - *args: The tensors to apply the function to. These tensors should be + Parameters + ---------- + f: Callable[..., array] + the function to apply elementwise. + args: array + the arrays to apply the function to. These arrays should be compatible for broadcasting. - Returns: - - LazyTensor: The tensor, `out`, of results from applying `f` elementwise to - the input tensors. After broadcasting the arguments to the same shape, for - each index `i`, `out[*i] = f(args[0][*i], args[1][*i], ...)`. + Returns + ------- + out: array + the array, ``out``, of results from applying ``f`` elementwise to + the input arrays. After broadcasting the arguments to the same shape, for + each index ``i``, ``out[*i] = f(args[0][*i], args[1][*i], ...)``. """ From de22d62b0629977c37606c5c16f8696bbf041e95 Mon Sep 17 00:00:00 2001 From: Willow Ahrens Date: Thu, 29 Oct 2026 12:25:14 -0400 Subject: [PATCH 3/6] fix --- src/array_api_stubs/_draft/einop.py | 60 +++++++++++++++++++++++++++++ 1 file changed, 60 insertions(+) create mode 100644 src/array_api_stubs/_draft/einop.py diff --git a/src/array_api_stubs/_draft/einop.py b/src/array_api_stubs/_draft/einop.py new file mode 100644 index 000000000..db433cf3e --- /dev/null +++ b/src/array_api_stubs/_draft/einop.py @@ -0,0 +1,60 @@ +__all__ = [ + "fft", + "ifft", + "fftn", + "ifftn", + "rfft", + "irfft", + "rfftn", + "irfftn", + "hfft", + "ihfft", + "fftfreq", + "rfftfreq", + "fftshift", + "ifftshift", +] + +from ._types import Tuple, Union, Sequence, array, Optional, Literal, dtype, device + +def einop(xp, prgm, **kwargs): + """Execute an einsum-like expression using the specified array framework. + + This function parses and executes einsum-like expressions with extended syntax + that supports various operations beyond traditional Einstein summation notation. + + Args: + xp: Array framework module (e.g., numpy, cupy, or other array library) + that provides the underlying array operations. + prgm (str): Einsum program string specifying the computation. The syntax + supports: + - Assignment: "C[i,j] = A[i,j] + B[j,i]" + - Increment: "C[i,j] += A[i,k] * B[k,j]" + - Reductions: "C[i] += A[i,j]", "C[i] max= A[i,j]", "C[i] &= A[i,j]" + - Arithmetic operations: +, -, *, /, //, %, ** + - Comparison operations: ==, !=, <, <=, >, >= + - Logical operations: and, or, not + - Bitwise operations: &, |, ^, <<, >> + - Function calls and complex expressions with parentheses + - Mathematical functions: abs, sqrt, exp, log, sin, cos, tan, etc. + - Literal values: integers, floats, booleans, and complex numbers + - Python operator precedence and parentheses for grouping + **kwargs: Named arrays referenced in the einsum expression. The keys + should match the tensor names used in the program string. + + Returns: + The result array from executing the einsum expression. + + Examples: + >>> import numpy as np + >>> A = np.random.rand(3, 4) + >>> B = np.random.rand(4, 3) + >>> # Matrix addition with transpose + >>> C = einsum(np, "C[i,j] = A[i,j] + B[j,i]", A=A, B=B) + >>> # Matrix multiplication + >>> D = einsum(np, "D[i,j] += A[i,k] * B[k,j]", A=A, B=B) + >>> # Min-Plus multiplication with shift + >>> E = einsum(np, "E[i] min= A[i,k] + D[k,j] << 1", A=A, D=D) + """ + prgm = parse_einop(prgm) + return prgm.run(xp, kwargs) \ No newline at end of file From a816cdb294f8da7d93e0f4dc2986ac1998be372d Mon Sep 17 00:00:00 2001 From: Willow Ahrens Date: Thu, 29 Oct 2026 12:37:48 -0400 Subject: [PATCH 4/6] stuff --- spec/draft/API_specification/einop.rst | 21 ++++++ spec/draft/API_specification/index.rst | 1 + src/array_api_stubs/_draft/__init__.py | 1 + src/array_api_stubs/_draft/einop.py | 99 ++++++++++++-------------- 4 files changed, 67 insertions(+), 55 deletions(-) create mode 100644 spec/draft/API_specification/einop.rst diff --git a/spec/draft/API_specification/einop.rst b/spec/draft/API_specification/einop.rst new file mode 100644 index 000000000..526315c0e --- /dev/null +++ b/spec/draft/API_specification/einop.rst @@ -0,0 +1,21 @@ +Einop +===== + + Array API specification for einop. + +A conforming implementation of the array API standard must provide and support the following functions. + + +Objects in API +-------------- + +.. currentmodule:: array_api + +.. + NOTE: please keep the functions in alphabetical order + +.. autosummary:: + :toctree: generated + :template: method.rst + + einop diff --git a/spec/draft/API_specification/index.rst b/spec/draft/API_specification/index.rst index 190573bdd..503436e09 100644 --- a/spec/draft/API_specification/index.rst +++ b/spec/draft/API_specification/index.rst @@ -25,6 +25,7 @@ A conforming implementation of the array API standard must provide and support t creation_functions data_type_functions data_types + einop elementwise_functions function_and_method_signatures higher_order diff --git a/src/array_api_stubs/_draft/__init__.py b/src/array_api_stubs/_draft/__init__.py index 78a47f1ba..c9a6685f6 100644 --- a/src/array_api_stubs/_draft/__init__.py +++ b/src/array_api_stubs/_draft/__init__.py @@ -5,6 +5,7 @@ from .creation_functions import * from .data_type_functions import * from . import data_types as dtype +from .einop import * from .elementwise_functions import * from .higher_order import * from .indexing_functions import * diff --git a/src/array_api_stubs/_draft/einop.py b/src/array_api_stubs/_draft/einop.py index db433cf3e..2c5efd531 100644 --- a/src/array_api_stubs/_draft/einop.py +++ b/src/array_api_stubs/_draft/einop.py @@ -1,60 +1,49 @@ -__all__ = [ - "fft", - "ifft", - "fftn", - "ifftn", - "rfft", - "irfft", - "rfftn", - "irfftn", - "hfft", - "ihfft", - "fftfreq", - "rfftfreq", - "fftshift", - "ifftshift", -] - -from ._types import Tuple, Union, Sequence, array, Optional, Literal, dtype, device - -def einop(xp, prgm, **kwargs): - """Execute an einsum-like expression using the specified array framework. +__all__ = ["einop"] + +from ._types import array + + +def einop(prgm: str, /, **kwargs: array) -> array: + """ + Executes an einsum-like expression. This function parses and executes einsum-like expressions with extended syntax that supports various operations beyond traditional Einstein summation notation. - Args: - xp: Array framework module (e.g., numpy, cupy, or other array library) - that provides the underlying array operations. - prgm (str): Einsum program string specifying the computation. The syntax - supports: - - Assignment: "C[i,j] = A[i,j] + B[j,i]" - - Increment: "C[i,j] += A[i,k] * B[k,j]" - - Reductions: "C[i] += A[i,j]", "C[i] max= A[i,j]", "C[i] &= A[i,j]" - - Arithmetic operations: +, -, *, /, //, %, ** - - Comparison operations: ==, !=, <, <=, >, >= - - Logical operations: and, or, not - - Bitwise operations: &, |, ^, <<, >> - - Function calls and complex expressions with parentheses - - Mathematical functions: abs, sqrt, exp, log, sin, cos, tan, etc. - - Literal values: integers, floats, booleans, and complex numbers - - Python operator precedence and parentheses for grouping - **kwargs: Named arrays referenced in the einsum expression. The keys - should match the tensor names used in the program string. - - Returns: - The result array from executing the einsum expression. - - Examples: - >>> import numpy as np - >>> A = np.random.rand(3, 4) - >>> B = np.random.rand(4, 3) - >>> # Matrix addition with transpose - >>> C = einsum(np, "C[i,j] = A[i,j] + B[j,i]", A=A, B=B) - >>> # Matrix multiplication - >>> D = einsum(np, "D[i,j] += A[i,k] * B[k,j]", A=A, B=B) - >>> # Min-Plus multiplication with shift - >>> E = einsum(np, "E[i] min= A[i,k] + D[k,j] << 1", A=A, D=D) + Parameters + ---------- + prgm: str + einsum program string specifying the computation. The syntax **must** support: + + - Assignment: ``"C[i,j] = A[i,j] + B[j,i]"`` + - Increment: ``"C[i,j] += A[i,k] * B[k,j]"`` + - Reductions: ``"C[i] += A[i,j]"``, ``"C[i] max= A[i,j]"``, ``"C[i] &= A[i,j]"`` + - Arithmetic operations: ``+``, ``-``, ``*``, ``/``, ``//``, ``%``, ``**`` + - Comparison operations: ``==``, ``!=``, ``<``, ``<=``, ``>``, ``>=`` + - Logical operations: ``and``, ``or``, ``not`` + - Bitwise operations: ``&``, ``|``, ``^``, ``<<``, ``>>`` + - Function calls and complex expressions with parentheses + - Mathematical functions: ``abs``, ``sqrt``, ``exp``, ``log``, ``sin``, ``cos``, ``tan``, etc. + - Literal values: integers, floats, booleans, and complex numbers + - Python operator precedence and parentheses for grouping + kwargs: array + named arrays referenced in the einsum expression. The keys **should** match the + array names used in ``prgm``. + + Returns + ------- + out: array + the result array from executing the einsum expression. + + Notes + ----- + + .. code-block:: + + # Matrix addition with transpose + C = einop("C[i,j] = A[i,j] + B[j,i]", A=A, B=B) + # Matrix multiplication + D = einop("D[i,j] += A[i,k] * B[k,j]", A=A, B=B) + # Min-Plus multiplication with shift + E = einop("E[i] min= A[i,k] + D[k,j] << 1", A=A, D=D) """ - prgm = parse_einop(prgm) - return prgm.run(xp, kwargs) \ No newline at end of file From 5f9cf08d6bc865af5e25398d1f9041894767ce65 Mon Sep 17 00:00:00 2001 From: Willow Ahrens Date: Thu, 29 Oct 2026 12:38:18 -0400 Subject: [PATCH 5/6] helpful ops --- src/array_api_stubs/_draft/einop.py | 116 ++++++++++++++++++++++++++++ 1 file changed, 116 insertions(+) diff --git a/src/array_api_stubs/_draft/einop.py b/src/array_api_stubs/_draft/einop.py index 2c5efd531..ef9d0447b 100644 --- a/src/array_api_stubs/_draft/einop.py +++ b/src/array_api_stubs/_draft/einop.py @@ -2,6 +2,122 @@ from ._types import array +nary_ops = { + "+": "add", + "add": "add", + "-": "subtract", + "sub": "subtract", + "subtract": "subtract", + "*": "multiply", + "mul": "multiply", + "multiply": "multiply", + "/": "divide", + "div": "divide", + "divide": "divide", + "//": "floor_divide", + "fld": "floor_divide", + "floor_divide": "floor_divide", + "%": "remainder", + "mod": "remainder", + "remainder": "remainder", + "**": "power", + "pow": "power", + "power": "power", + "==": "equal", + "eq": "equal", + "equal": "equal", + "!=": "not_equal", + "ne": "not_equal", + "not_equal": "not_equal", + "<": "less", + "lt": "less", + "less": "less", + "<=": "less_equal", + "le": "less_equal", + "less_equal": "less_equal", + ">": "greater", + "gt": "greater", + "greater": "greater", + ">=": "greater_equal", + "ge": "greater_equal", + "greater_equal": "greater_equal", + "&": "bitwise_and", + "bitwise_and": "bitwise_and", + "|": "bitwise_or", + "bitwise_or": "bitwise_or", + "^": "bitwise_xor", + "bitwise_xor": "bitwise_xor", + "<<": "bitwise_left_shift", + "lshift": "bitwise_left_shift", + "bitwise_left_shift": "bitwise_left_shift", + ">>": "bitwise_right_shift", + "rshift": "bitwise_right_shift", + "bitwise_right_shift": "bitwise_right_shift", + "and": "logical_and", + "or": "logical_or", + "not": "logical_not", + "min": "minimum", + "max": "maximum", + "logaddexp": "logaddexp", +} + + +unary_ops = { + "+": "positive", + "pos": "positive", + "positive": "positive", + "-": "negative", + "neg": "negative", + "negative": "negative", + "~": "bitwise_invert", + "invert": "bitwise_invert", + "bitwise_invert": "bitwise_invert", + "not": "logical_not", + "logical_not": "logical_not", + "abs": "absolute", + "absolute": "absolute", + "sqrt": "sqrt", + "exp": "exp", + "log": "log", + "log1p": "log1p", + "log10": "log10", + "log2": "log2", + "sin": "sin", + "cos": "cos", + "tan": "tan", + "sinh": "sinh", + "cosh": "cosh", + "tanh": "tanh", + "asin": "arcsin", + "acos": "arccos", + "atan": "arctan", + "asinh": "arcsinh", + "acosh": "arccosh", + "atanh": "arctanh", +} + + +reduction_ops = { + "+": "sum", + "add": "sum", + "sum": "sum", + "*": "prod", + "mul": "prod", + "prod": "prod", + "and": "all", + "or": "any", + "min": "min", + "max": "max", + "argmin": "argmin", + "argmax": "argmax", + "mean": "mean", + "std": "std", + "var": "var", + "count_nonzero": "count_nonzero", + # "&": "bitwise_and", + # "|": "bitwise_or", + # "^": "bitwise_xor", +} def einop(prgm: str, /, **kwargs: array) -> array: """ From 44c0e40cbaa13d33a138755322a43a9d06f902df Mon Sep 17 00:00:00 2001 From: Willow Ahrens Date: Thu, 29 Oct 2026 13:18:24 -0400 Subject: [PATCH 6/6] fix --- src/array_api_stubs/_draft/einop.py | 118 +--------------------------- 1 file changed, 2 insertions(+), 116 deletions(-) diff --git a/src/array_api_stubs/_draft/einop.py b/src/array_api_stubs/_draft/einop.py index ef9d0447b..29522c74d 100644 --- a/src/array_api_stubs/_draft/einop.py +++ b/src/array_api_stubs/_draft/einop.py @@ -2,122 +2,6 @@ from ._types import array -nary_ops = { - "+": "add", - "add": "add", - "-": "subtract", - "sub": "subtract", - "subtract": "subtract", - "*": "multiply", - "mul": "multiply", - "multiply": "multiply", - "/": "divide", - "div": "divide", - "divide": "divide", - "//": "floor_divide", - "fld": "floor_divide", - "floor_divide": "floor_divide", - "%": "remainder", - "mod": "remainder", - "remainder": "remainder", - "**": "power", - "pow": "power", - "power": "power", - "==": "equal", - "eq": "equal", - "equal": "equal", - "!=": "not_equal", - "ne": "not_equal", - "not_equal": "not_equal", - "<": "less", - "lt": "less", - "less": "less", - "<=": "less_equal", - "le": "less_equal", - "less_equal": "less_equal", - ">": "greater", - "gt": "greater", - "greater": "greater", - ">=": "greater_equal", - "ge": "greater_equal", - "greater_equal": "greater_equal", - "&": "bitwise_and", - "bitwise_and": "bitwise_and", - "|": "bitwise_or", - "bitwise_or": "bitwise_or", - "^": "bitwise_xor", - "bitwise_xor": "bitwise_xor", - "<<": "bitwise_left_shift", - "lshift": "bitwise_left_shift", - "bitwise_left_shift": "bitwise_left_shift", - ">>": "bitwise_right_shift", - "rshift": "bitwise_right_shift", - "bitwise_right_shift": "bitwise_right_shift", - "and": "logical_and", - "or": "logical_or", - "not": "logical_not", - "min": "minimum", - "max": "maximum", - "logaddexp": "logaddexp", -} - - -unary_ops = { - "+": "positive", - "pos": "positive", - "positive": "positive", - "-": "negative", - "neg": "negative", - "negative": "negative", - "~": "bitwise_invert", - "invert": "bitwise_invert", - "bitwise_invert": "bitwise_invert", - "not": "logical_not", - "logical_not": "logical_not", - "abs": "absolute", - "absolute": "absolute", - "sqrt": "sqrt", - "exp": "exp", - "log": "log", - "log1p": "log1p", - "log10": "log10", - "log2": "log2", - "sin": "sin", - "cos": "cos", - "tan": "tan", - "sinh": "sinh", - "cosh": "cosh", - "tanh": "tanh", - "asin": "arcsin", - "acos": "arccos", - "atan": "arctan", - "asinh": "arcsinh", - "acosh": "arccosh", - "atanh": "arctanh", -} - - -reduction_ops = { - "+": "sum", - "add": "sum", - "sum": "sum", - "*": "prod", - "mul": "prod", - "prod": "prod", - "and": "all", - "or": "any", - "min": "min", - "max": "max", - "argmin": "argmin", - "argmax": "argmax", - "mean": "mean", - "std": "std", - "var": "var", - "count_nonzero": "count_nonzero", - # "&": "bitwise_and", - # "|": "bitwise_or", - # "^": "bitwise_xor", -} def einop(prgm: str, /, **kwargs: array) -> array: """ @@ -154,6 +38,8 @@ def einop(prgm: str, /, **kwargs: array) -> array: Notes ----- + - Operations **must** be evaluated as with :func:`~array_api.elementwise`, and reductions as with :func:`~array_api.reduce`. + .. code-block:: # Matrix addition with transpose