Isaac Breen
← Projects and writing

einexpr

An experimental Python library for named tensor dimensions and Einstein-style array programming.

Period
2022–2024
Status
Experimental · archived
Project
GitHub repository
Package
Package page

einexpr is a Python framework I wrote for working with named tensor dimensions. You can write the dimension names directly on arrays and use them to align, contract, rename and group dimensions.

import einexpr as ei

a = ei.array([1, 2])
b = ei.array([3, 4])
X = ei.array([[1, 2], [3, 4]])
Y = ei.array([[5, 6], [7, 8]])

dot = (a['i'] * b['i'])['']
outer = (a['i'] * b['j'])['i j']

product = (X['i j'] * Y['j k'])['i k']

The names determine how dimensions line up and which ones are contracted. In the matrix product, j appears in both inputs and disappears from the result. The expression does not depend on j occupying a particular axis.

Multi-head attention

Here is multi-head attention written both ways. PyTorch uses reshapes, transposes and axis numbers. einexpr names the dimensions on each array instead. t->s renames the key/value position before the contraction, and the result names say which dimensions remain.

The package also supported grouped dimensions. For example, context['b t (h d)'] joins the head and per-head feature dimensions while keeping the grouping explicit in the expression.

Projections
qbhtd=xbtcWchdQq_{bhtd}=x_{btc}W^Q_{chd}
kbhsd=xbscWchdKk_{bhsd}=x_{bsc}W^K_{chd}
vbhsd=xbscWchdVv_{bhsd}=x_{bsc}W^V_{chd}
PyTorch 
q = (x @ w_q.reshape(C, H * D)).reshape(B, T, H, D).transpose(1, 2)
k = (x @ w_k.reshape(C, H * D)).reshape(B, T, H, D).transpose(1, 2)
v = (x @ w_v.reshape(C, H * D)).reshape(B, T, H, D).transpose(1, 2)
einexpr 
q = (x['b t c'] * w_q['c h d'])['b h t d']
k = (x['b t c'] * w_k['c h d'])['b h t d']
v = (x['b t c'] * w_v['c h d'])['b h t d']
Scores
rbhts=qbhtdkbhsdDr_{bhts}=\frac{q_{bhtd}k_{bhsd}}{\sqrt D}
PyTorch 
scores = (q @ k.transpose(-2, -1)) / math.sqrt(D)
einexpr 
scores = (q['b h t d'] * k['b h t->s d'])['b h t s'] / math.sqrt(D)
Softmax
mbht=max⁡srbhtsm_{bht}=\max_s r_{bhts}
a~bhts=exp⁡(rbhts−mbht)\tilde a_{bhts}=\exp(r_{bhts}-m_{bht})
abhts=a~bhts∑s′a~bhts′a_{bhts}=\frac{\tilde a_{bhts}}{\sum_{s'}\tilde a_{bhts'}}
PyTorch 
max_scores = scores.amax(dim=-1, keepdim=True)
weights = torch.exp(scores - max_scores)
weights = weights / weights.sum(dim=-1, keepdim=True)
einexpr 
max_scores = ei.max(scores, axis='s')
weights = ei.exp(scores - max_scores)
weights = weights / ei.sum(weights, axis='s')
Output
pbthd=abhtsvbhsdp_{bthd}=a_{bhts}v_{bhsd}
ybtc=pbthdWhdcOy_{btc}=p_{bthd}W^O_{hdc}
PyTorch 
context = (weights @ v).transpose(1, 2).reshape(B, T, H * D)
y = context @ w_o.reshape(H * D, C)
einexpr 
context = (weights['b h t s'] * v['b h t->s d'])['b t h d']
y = (context['b t h d'] * w_o['h d c'])['b t c']

Coda

A whole Transformer block, for fun

The EinsteinTensors notebooks attempted something like this in 2021, but the equations were unfinished. This is a corrected reconstruction: a post-layer-normalisation causal Transformer block, with dropout omitted.

The mathematics and code below are each meant to stand on their own. Hover—or keyboard-focus—any equation or code line to mark its counterpart in the other block.

einexpr

# Attention
q = (x['b t c'] * w_q['c h d'])['b h t d'] + b_q['h d']
k = (x['b t c'] * w_k['c h d'])['b h t d'] + b_k['h d']
v = (x['b t c'] * w_v['c h d'])['b h t d'] + b_v['h d']

scores = (
    q['b h t d'] * k['b h t->s d']
)['b h t s'] / math.sqrt(D)
scores = scores + mask['t s']

weights = ei.exp(scores - ei.max(scores, axis='s'))
weights = weights / ei.sum(weights, axis='s')

context = (
    weights['b h t s'] * v['b h t->s d']
)['b t h d']

attention = (
    context['b t h d'] * w_o['h d c']
)['b t c'] + b_o['c']

# Residual and layer norm
u1 = x + attention
mean_u = ei.mean(u1, axis='c')
var_u = ei.mean((u1 - mean_u) ** 2, axis='c')
u0 = (
    gamma_u['c'] * (u1 - mean_u) / ei.sqrt(var_u + eps)
    + beta_u['c']
)

# Feed-forward network
hidden = gelu(
    (u0['b t c'] * w_1['c m'])['b t m'] + b_1['m']
)
ff = (hidden['b t m'] * w_2['m c'])['b t c'] + b_2['c']

# Residual and layer norm
z1 = u0 + ff
mean_z = ei.mean(z1, axis='c')
var_z = ei.mean((z1 - mean_z) ** 2, axis='c')
z0 = (
    gamma_z['c'] * (z1 - mean_z) / ei.sqrt(var_z + eps)
    + beta_z['c']
)
qbhtd=xbtcWchdQ+bhdQq_{bhtd}=x_{btc}W^Q_{chd}+b^Q_{hd}
kbhsd=xbscWchdK+bhdKk_{bhsd}=x_{bsc}W^K_{chd}+b^K_{hd}
vbhsd=xbscWchdV+bhdVv_{bhsd}=x_{bsc}W^V_{chd}+b^V_{hd}
rbhts=qbhtdkbhsdD+Mtsr_{bhts}=\frac{q_{bhtd}k_{bhsd}}{\sqrt D}+M_{ts}
abhts=softmax⁡s ⁣(rbhts)a_{bhts}=\operatorname{softmax}_{s}\!\left(r_{bhts}\right)
pbthd=abhtsvbhsdp_{bthd}=a_{bhts}v_{bhsd}
obtc=pbthdWhdcO+bcOo_{btc}=p_{bthd}W^O_{hdc}+b^O_c
ubtc1=xbtc+obtcu^1_{btc}=x_{btc}+o_{btc}
μbtu=1C∑cubtc1\mu^u_{bt}=\frac{1}{C}\sum_c u^1_{btc}
(σbtu)2=1C∑c(ubtc1−μbtu)2(\sigma^u_{bt})^2=\frac{1}{C}\sum_c\left(u^1_{btc}-\mu^u_{bt}\right)^2
ubtc0=γcuubtc1−μbtu(σbtu)2+ε+βcuu^0_{btc}=\gamma^u_c\frac{u^1_{btc}-\mu^u_{bt}}{\sqrt{(\sigma^u_{bt})^2+\varepsilon}}+\beta^u_c
fbtm=GELU⁡ ⁣(ubtc0Wcm1+bm1)f_{btm}=\operatorname{GELU}\!\left(u^0_{btc}W^1_{cm}+b^1_m\right)
gbtc=fbtmWmc2+bc2g_{btc}=f_{btm}W^2_{mc}+b^2_c
zbtc1=ubtc0+gbtcz^1_{btc}=u^0_{btc}+g_{btc}
μbtz=1C∑czbtc1\mu^z_{bt}=\frac{1}{C}\sum_c z^1_{btc}
(σbtz)2=1C∑c(zbtc1−μbtz)2(\sigma^z_{bt})^2=\frac{1}{C}\sum_c\left(z^1_{btc}-\mu^z_{bt}\right)^2
zbtc0=γczzbtc1−μbtz(σbtz)2+ε+βczz^0_{btc}=\gamma^z_c\frac{z^1_{btc}-\mu^z_{bt}}{\sqrt{(\sigma^z_{bt})^2+\varepsilon}}+\beta^z_c
Parameter initialisation and setup Show the code hidden behind the main block 
import math
import numpy as np
import einexpr as ei

B, T, C = 2, 8, 64
H, D, M = 4, 16, 256
eps = 1e-5
rng = np.random.default_rng(0)

def tensor(shape, dims):
    return ei.asarray(rng.normal(size=shape), dims=dims)

def gelu(x):
    return 0.5 * x * (
        1.0
        + ei.tanh(
            math.sqrt(2.0 / math.pi)
            * (x + 0.044715 * x**3)
        )
    )

x = tensor((B, T, C), 'b t c')

w_q = tensor((C, H, D), 'c h d')
b_q = tensor((H, D), 'h d')
w_k = tensor((C, H, D), 'c h d')
b_k = tensor((H, D), 'h d')
w_v = tensor((C, H, D), 'c h d')
b_v = tensor((H, D), 'h d')
w_o = tensor((H, D, C), 'h d c')
b_o = tensor((C,), 'c')

gamma_u = ei.ones((C,), dims='c')
beta_u = ei.zeros((C,), dims='c')
w_1 = tensor((C, M), 'c m')
b_1 = tensor((M,), 'm')
w_2 = tensor((M, C), 'm c')
b_2 = tensor((C,), 'c')
gamma_z = ei.ones((C,), dims='c')
beta_z = ei.zeros((C,), dims='c')

mask = np.where(
    np.arange(T)[None, :] <= np.arange(T)[:, None],
    0.0,
    -np.inf,
)
mask = ei.asarray(mask, dims='t s')

The displayed program was run against the archived package and compared with an independent NumPy implementation of the same block; the maximum difference was 1.27e-14.

Started Jun 8, 2022 · Updated Aug 11, 2026