Extended Einsum

Get started View on GitHub

Extended Einsum is a small tensor language and compiler IR. It keeps contractions, intermediate results, layout operations, and nonlinearities visible long enough to optimize the whole program—then lowers it to PyTorch, JAX, or NumPy.

Beyond one einsum

Compose contractions with softmax, exponentials, arithmetic, stacking, slicing, selection, and routing.

Whole-program optimization

Replan contraction paths, fold matching operations, and arrange folded values for their consumers.

Automatic stability

Choose log-space or scaled evaluation without rewriting the expression that defines your model.

Backend independent

Use the included PyTorch, JAX, and NumPy backends or implement the compact backend protocol.

Quick start

import torch
import extended_einsum.interface as xe

x = xe.array(torch.rand(32, 16) + 0.1)
w = xe.array(torch.rand(16, 8) + 0.1)

hidden = xe.softmax(x, axis=1)
result = xe.einsum("bi,io->bo", hidden, w)
output = result.materialize(stability_mode="scaled_sum")

print(output.backend_array.shape)  # torch.Size([32, 8])

import numpy as np
import extended_einsum.interface as xe

x = xe.array(np.random.rand(32, 16) + 0.1)
w = xe.array(np.random.rand(16, 8) + 0.1)
result = xe.einsum("bi,io->bo", xe.softmax(x, axis=1), w)

output = result.materialize(stability_mode="unstable")
print(output.backend_array.shape)  # (32, 8)

import jax.numpy as jnp
import extended_einsum.interface as xe

x = xe.array(jnp.ones((32, 16)))
w = xe.array(jnp.ones((16, 8)))
result = xe.einsum("bi,io->bo", xe.softmax(x, axis=1), w)

output = result.materialize(stability_mode="scaled_sum")
print(output.backend_array.shape)  # (32, 8)

Choose a path

New to the package? Start with installation and your first expression.
Training a model? Use the minimal PyTorch training setup.
Working with deep products of positive values? Read automatic numerical stability.
Integrating another array framework? Follow write a backend.
Looking up a signature? Go directly to the interface reference.

The compilation pipeline

TensorExpression
    ↓ extract_program
Extended Einsum SSA program
    ↓ folding and contraction-path passes
Optimized program
    ↓ stability-aware translation
BackendProgram
    ↓ PyTorch / JAX / NumPy compiler
Native backend result

The frontend records what the tensor program computes. Optimization decides how to structure it. Stability translation chooses how positive values are represented. The backend finally supplies array primitives and optional compilation.