einx: Universal Notation for Tensor Operations in Python

einx provides a universal interface for tensor operations across frameworks like Numpy, PyTorch, Jax, and Tensorflow using a specific string notation. Every operation follows this pattern:

outputs... = einx.{elementary_operation}("{vectorization}", inputs...)

The vectorization string serves three purposes:

1. Full Operation Signature: Defines the relationship between all input and output dimensions. Inputs and outputs are separated by ->, and tensors are separated by ,.

   • Example: "[c d] a, b -> a [e] b" accepts shapes (c, d, a) and (b), returning (a, e, b).

2. Elementary Operation Signature: Sub-expressions in brackets [] define the dimensions that the core operation (e.g., sum, dot, add) actually operates on. Axes outside brackets are treated as vectorized (loop) axes.

   • Example: In "[c d] a, b -> a [e] b", the elementary operation acts on shapes (c, d) and () and returns (e).

3. Vectorization: Axes not in brackets are handled via vectorization (analogous to nested loops).

   • Example: "[c d] a, b -> a [e] b" is equivalent to:

     for a in range(...):
         for b in range(...):
             z[a, :, b] = elementary_operation(x[:, :, a], y[b])

Note: einx uses backend-optimized functions rather than literal Python loops.

While einops.pack and einops.unpack use a special * notation for concatenation, einx handles these operations using its standard notation via einx.id. This allows for more expressive operations where input and output shapes do not align, and provides full documentation of all axes in the expression.

Common concatenation patterns in einx:

• Concatenate along a dimension: einx.id("a b c1, a b c2 -> a b (c1 + c2)", x, y)
• Append a constant to a dimension: einx.id("b h w c1, -> b h w (c1 + 1)", img, 42.0)
• Create mesh-grids: einx.id("x, y -> x y (1 + 1)", x, y)

# Concatenation of tensors along the third dimension
z = einx.id("a b c1, a b c2 -> a b (c1 + c2)", x, y)

# Appending a number to the channel dimension of an image
img_out = einx.id("b h w c1, -> b h w (c1 + 1)", img, 42.0)

# Creating mesh-grids (similar to np.meshgrid)
x = np.arange(64)
y = np.arange(48)
grid = einx.id("x, y -> x y (1 + 1)", x, y)

An ellipsis (...) in an einx expression represents multiple axes jointly. It is placed immediately after a sub-expression to indicate that the sub-expression is repeated zero or more times. The number of repetitions is automatically inferred from the input tensor shapes and any provided constraints.

Key Behaviors:

• Inference: The number of axes covered by the ellipsis is determined by the input dimensionality.
• Consistency: If the same axis name is used with an ellipsis multiple times (e.g., s...), the number of repetitions must match across all occurrences. A mismatch raises an exception.
• Composition: Ellipses can be applied to complex sub-expressions, including those with brackets [] or axis compositions.
• Anonymous Ellipses: You can use ... without a preceding expression (e.g., ..., ... -> ...). In this case, einx generates a new, unique axis name for all occurrences.
• Constraints: You can provide additional constraints for axes expanded by ellipses as either a tuple of integers (matching the repetition count) or a single integer that applies to all repetitions.

import numpy as np
import einx

# 1. Basic ellipsis expansion
x = np.random.randn(10, 20, 30, 40)
y = einx.sum("s... [c] -> s...", x)
# expands to: einx.sum("s1 s2 s3 [c] -> s1 s2 s3", x)

# 2. Multiple ellipses with different sub-expressions
x = np.random.randn(10, 20)
y = np.random.randn(30, 40)
z = einx.add("a..., b... -> a... b...", x, y)
# expands to: einx.add("a1 a2, b1 b2 -> a1 a2 b1 b2", x, y)

# 3. Ellipses with constraints
# Using a tuple for specific sizes per repetition
y = einx.id("a -> a b...", x, b=(5, 6))
# expands to: einx.id("a -> a b1 b2", x, b1=5, b2=6)

# Using a single integer for all repetitions
y = einx.id("(a b)... -> a... b...", x, b=2)
# expands to: einx.id("(a1 b1) (a2 b2) -> a1 a2 b1 b2", x, b1=2, b2=2)

# 4. Anonymous ellipses
z = einx.add("..., ... -> ...", x, y)

A core advantage of einx is that it decouples the representation of an elementary operation from its vectorization (the shape of the tensors). This allows you to change the input/output shapes by simply updating the notation string, whereas traditional libraries (like NumPy or PyTorch) often require switching to different functions or adding manual rearrange steps when shapes change.

Changing the Shape

In einx, you can adapt to different indexing or broadcasting requirements by modifying the notation string while keeping the same entry-point function. For example, when moving from 1D indexing to more complex multidimensional indexing, einx handles the necessary rearrangements internally.

# 1D indexing
einx.get_at("[x] a, b -> b a", x, y)

# 2D indexing (einx handles the complexity that would require manual rearrange in PyTorch)
einx.get_at("[x y] b, c b a [2] -> c b a", x, y)

Tensor notation can be categorized into two main styles:

1. Classical (Positional) Notation: Dimensions are identified by their position in the shape (e.g., the 1st, 2nd, or 3rd dimension). This often requires comments or strict conventions (like "channels-last") to prevent errors.
2. Named Tensor Notation: Dimensions are annotated with symbolic names (e.g., batch, feature, time). This makes code self-documenting and allows operations to target specific axes by name rather than index.

einx is compatible with both. An einx operation can match its string expression against either the positional shape of a classical tensor or the symbolic axis names of a named tensor.

# Classical, positional-style tensor
x = create_tensor((32, 128, 128))
y = sum(x, axis=2)

# Named tensor
x = create_tensor({"batch": 32, "feature": 128, "time": 128})
y = sum(x, axis="time")

# einx works with both by matching names or positions
y = einx.sum("batch feature [time]", x)