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/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..503436e09 100644 --- a/spec/draft/API_specification/index.rst +++ b/spec/draft/API_specification/index.rst @@ -25,8 +25,10 @@ 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 indexing indexing_functions inspection diff --git a/src/array_api_stubs/_draft/__init__.py b/src/array_api_stubs/_draft/__init__.py index 537ea8f85..c9a6685f6 100644 --- a/src/array_api_stubs/_draft/__init__.py +++ b/src/array_api_stubs/_draft/__init__.py @@ -5,7 +5,9 @@ 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 * 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/einop.py b/src/array_api_stubs/_draft/einop.py new file mode 100644 index 000000000..29522c74d --- /dev/null +++ b/src/array_api_stubs/_draft/einop.py @@ -0,0 +1,51 @@ +__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. + + 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 + ----- + + - Operations **must** be evaluated as with :func:`~array_api.elementwise`, and reductions as with :func:`~array_api.reduce`. + + .. 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) + """ diff --git a/src/array_api_stubs/_draft/higher_order.py b/src/array_api_stubs/_draft/higher_order.py new file mode 100644 index 000000000..93054e9f1 --- /dev/null +++ b/src/array_api_stubs/_draft/higher_order.py @@ -0,0 +1,199 @@ +__all__ = ["reduce", "elementwise"] + +from ._types import Callable, Optional, Tuple, Union, array, dtype + + +binary_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", + "**": "pow", + "pow": "pow", + "power": "pow", + "==": "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", + "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": "abs", + "absolute": "abs", + "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": "asin", + "acos": "acos", + "atan": "atan", + "asinh": "asinh", + "acosh": "acosh", + "atanh": "atanh", +} + + +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: Callable[[array, array], array], + x: array, + /, + *, + axis: Optional[Union[int, Tuple[int, ...]]] = None, + dtype: Optional[dtype] = None, + keepdims: bool = False, + 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``. + + 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, ...]]] + 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[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 + ``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: 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 arrays to + ensure they have compatible shapes. For example, ``elementwise(add, x, y)`` + is equivalent to ``x + y``. + + 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 + ------- + 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], ...)``. + """