From bc8d9a6a86931c84d82c27639e721bb9a28a2262 Mon Sep 17 00:00:00 2001 From: Willow Ahrens Date: Thu, 29 Oct 2026 12:58:29 -0400 Subject: [PATCH 1/2] 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/2] 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], ...)``. """