Wma/einsum

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

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

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

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 *

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,

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)
+    """

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], ...)``.
+    """