einx: one string notation for tensor operations across NumPy, PyTorch, JAX, TensorFlow and MLX

The problem einx attacks is shape bookkeeping, not arithmetic

Most tensor code is a negotiation with axis order. You have a tensor, you know its shape by heart, and then you call a framework function with a tuple of integers that only you can read. Swap that call to another framework and every one of those integer tuples has to be translated again, even though the intent, say a mean pool over the two spatial axes, never changed. The einx project exists to remove that translation step. Its description calls it a universal notation for tensor operations in Python, and the README frames it as a notation plus a library that offers one interface to Numpy, PyTorch, Jax, Tensorflow and MLX. The audience is therefore not a beginner learning to manipulate arrays. It is the engineer who maintains research code that has to run under more than one backend, where a hardcoded transpose is a portability bug waiting for the next experiment. einx moves the axis specification out of positional integers and into a string that names axes, so the same expression can be handed to any of the supported frameworks.

An einx expression is a loop, written once and compiled away

The design idea is simple to state and worth understanding precisely. Every einx operation combines an elementary operation with an expression that says how that operation is vectorized. The project documentation defines the meaning of the string by analogy with loop notation: the full operation yields the same output as if the elementary operation were invoked in an equivalent Python loop. Under that reading, every element of the first axis of x and every element of the second axis of y are combined separately, and the result is assembled in the new axis order. The important detail is that the loop is a specification of the result, not the implementation. Internally, operations are compiled to Python code snippets that call into the respective tensor framework, so there is no Python level iteration over elements at runtime. That distinction is what separates einx from a naive wrapper that loops in Python and would be far too slow for anything real. Axes in square brackets are grouped and flattened together, a plus sign inside a group such as "(q + k)" marks a split into two named halves, and a trailing comma in the output such as "b c, -> b (c + 1)" appends a size one axis.

graph=True shows you the code before it runs

Because einx compiles rather than interprets, it can hand you the generated snippet. Passing graph=True to any operation returns the Python source that would execute, which is the fastest way to settle a question about axis order. The project shows this for an addition with a transposed output, where the compiled body transposes and reshapes both inputs before calling np.add. The result is readable NumPy, and the same mechanism applies to the other backends, where a vmap based implementation looks different. Generated code is also the only place where you can see exactly how einx decided to flatten a group, so if a result looks wrong, printing the graph is the first step rather than a last resort. The downside is that the snippet is a debugging artifact rather than a stable interface, so code that depends on its exact formatting will age badly. The project documentation also points out that different backends may compile an operation to different implementations, for example a Numpy-like form, a vmap based form or an einsum form.

Installing einx and running a first reduction

einx is a pure Python package published on PyPI, and the quickstart gives one command for installation. The dependency list in pyproject.toml is short, numpy and sympy, with torch and mlx offered as optional extras. A first call looks like this.

pip install einx

The same reduction reads the same on every backend

The quickstart then creates a tensor and calls one einx operation on it. The string names the three axes of the input as a, b and c, and the bracket around b tells the library to reduce over that axis as a group. The operation returns a new array rather than modifying the input, and printing the shape is how you confirm that the axis you meant to collapse is the one that was collapsed.

import einx
import numpy as np

x = np.ones((10, 20, 30)) # Create some tensor (numpy/torch/jax/tensorflow/mlx/...)

y = einx.sum("a [b] c", x) # Call an einx operation

print(y.shape)

One grammar covers permutations, splits, pools and gathers

The breadth of the notation is best judged from the set of examples the README lists. The same expression grammar carries a permutation with flattening, a split into two named halves, an append of a size one axis, a reduction over a bracket group, a mean pool over the spatial axes, an outer sum, a matrix multiplication, and a gather at coordinates. The unifying idea is that each of these is an elementary operation plus a description of how it is vectorized, so nothing needs a bespoke keyword argument set. The bracket group is what lets a single expression express a two by two sum pool, since the group size can be given as a tuple.

z = einx.id("a (b c) -> (b a) c", x, b=2)             # Permute and (un)flatten axes
z = einx.id("b (q + k) -> b q, b k", x, q=2)          # Split
z = einx.id("b c, -> b (c + 1)", x, 42)               # Append number to each channel
z = einx.sum("a [b]", x)                              # Sum-reduction along second axis
z = einx.flip("... (g [c])", x, c=2)                  # Flip pairs of values along the last axis
z = einx.mean("b [...] c", x)                         # Spatial mean-pooling
z = einx.multiply("a..., b... -> (a b)...", x, y)     # Kronecker product
z = einx.sum("b (s [ds])... c", x, ds=(2, 2))         # Sum-pooling with 2x2 kernel
z = einx.add("a, b -> a b", x, y)                     # Outer sum
z = einx.dot("a [b], [b] c -> a c", x, y)             # Matrix multiplication
z = einx.get_at("b [h w] c, b i [2] -> b i c", x, y)  # Gather values at coordinates

A custom function can be adapted, at the cost of vectorization control

einx is not limited to the built in operation set. A plain Python function can be adapted to the notation with an adapter, and the project shows the pattern using einx.torch.adapt_with_vmap, which wraps the function so that torch.vmap performs the batching. The adapted function keeps the interface of a normal einx operation, including the graph=True inspection. The trade off is visible in the generated code: the wrapper asserts that the adapted function returns a tensor and that its shape matches what the expression expects, so a function returning a scalar or an unexpected shape fails with an assertion rather than a silent broadcast. The adapter list lives in its own documentation page, and the compile strategy depends on the backend, so an adapter that works under torch does not guarantee identical behaviour elsewhere.

# Define a custom operation
def foo(x, y):
    x = 2 * x
    z = x + torch.sum(y)
    return z

# Adapt the operation to einx notation
einfoo = einx.torch.adapt_with_vmap(foo)

# Invoke as einx operation
z = einfoo("a [c], b [c] -> a b [c]", x, y)

Where einx is the wrong tool

einx is a notation layer, and the limits of that choice show up in three places. First, the actual compute still happens in NumPy, PyTorch, JAX, Tensorflow or MLX, so you inherit the performance characteristics and the device placement rules of whichever backend you picked, with no batching layer of its own. Second, the notation describes shapes and axis names, not semantics, so an expression that is legal but ambiguous, such as a group of axes whose size does not divide evenly, has to be caught by the compiler rather than by the reader. Third, the backends are not identical, so the same expression can have different performance profiles across frameworks even though the output matches. The library is also classified as Development Status :: 3 - Alpha in its packaging metadata, which is a signal that the expression grammar and the operation set may still shift between minor versions. If your team wants one native implementation of a heavy kernel rather than one expression that dispatches to several, einx is the wrong abstraction.

Maintenance, packaging and the cost of an upgrade

The repository is not archived, and the last push was on 2026-09-12. The most recent tagged releases are v0.4.1 on 2026-03-05, v0.4.2 on 2026-03-14 and v0.4.3 on 2026-04-01, and the project version in pyproject.toml is 0.4.3, so the packaging metadata and the newest tag agree. Releases arrive as a run of small numbers rather than on a fixed calendar, so there is no monthly cadence to plan around and the changelog deserves a read before each bump. The license is MIT, with the license file shipped alongside the package, and the documentation lives at einx.readthedocs.io. Python 3.10 or newer is required. Type information is shipped, with py.typed in the package data and a separate test_typechecking directory at the repository root, so annotated code can be checked without a separate stub package.

Editorial conclusion

einx suits engineers who write model code that must run on more than one tensor framework and who are willing to read a generated graph to debug shapes. It does not suit teams that want one implementation of heavy numeric kernels, since the work still happens inside each framework, and it is a poor fit for code that must stay readable to people who do not already know loop notation. Before committing, verify three things in your own repository: that the backends you need appear in the supported list, that the compiled output for your most common operations matches what you expect from einsum, and that a released version is pinned rather than tracked from main.

Frequently asked questions

Official sources

1. fferflo/einx on GitHub
2. License: MIT
3. Project website
4. README
5. Releases