Graphical Design of Interpretable Architectures

Designing, implementing, and comparing interpretable architectures requires a formal language to represent them. The most common representations fall short in one of two ways. Symbolic equations give no global view of an architecture at a glance. Probabilistic graphical models and flowcharts do not describe actual tensor manipulations, thus hiding key insights and limiting reproducibility. To close this gap, we introduce a graphical notation for designing interpretable AI architectures, adapted from Penrose tensor notation. This graphical notation gives a global view of an architecture and maps one to one onto PyTorch einsum code. We first use this notation to describe architectures that are interpretable by construction, including concept bottlenecks, sparse probes, prototype networks, neural additive models, and mixtures of linear models. We then diagram the key architectural components of Steerling-8B, a frontier interpretable language model. The diagram yields global insights into the architecture (e.g., showing that Steerling is a residual model), a geometric interpretation of each individual operation, and a direct translation into 33 lines of PyTorch code.

Introduction

Frontier AI models manipulate high-dimensional objects called tensors. For this reason, to understand, design, or implement these models, we must think in high-dimensional terms. As frontier models are composed of many such operations, intuitive and formal representations of their tensor manipulations are key to understanding, comparing, and designing state-of-the-art architectures.

Common formal representations, such as symbolic equations, may limit global insights on the architecture at a glance. As an example, the expression

\[f(x_k,w_{kri},t_{krij}) = \sum_i \sigma(x_{k}w_{kri}) t_{krij}\]

hides that the tensor manipulations act on individual features $k$ of $x$ independently. We must carefully work through the whole expression to see this. This process requires time, effort, and it is intrinsically prone to errors. Graphical notations, such as probabilistic graphical models and flowcharts, are commonly used specifically to compensate the shortcomings of symbolic representations and convey immediate insights:

However, while these diagrams provide a high-level overview of the architecture, they do not specify how the architecture should be concretely implemented thus obfuscating important innovations and limiting reproducibility.

To solve this, we introduce a graphical notation, adapted from Penrose tensor notation, for designing interpretable AI architectures that enables global insights and maps one to one onto PyTorch einsum code. While prior work has used a similar notation to analyse language models and mechanistic interpretability, we use this notation to describe architectures that are interpretable by construction, including concept bottlenecks, sparse probes, prototype networks, neural additive models, and mixtures of linear models. As a case study, we diagram the full architecture of Steerling-8B, showing the advantages of such notation in practice.

Graphical Einstein-inspired notation in PyTorch

A well-known graphical notation for tensor operations, originating in physics, is Penrose or tensor-network notation . Recent work has adapted it to analyse language models and mechanistic interpretability . To our knowledge, no prior work has used it to analyse models that are interpretable by construction.

We first introduce the fragment of Penrose notation we need. We then use it to analyse the key operations behind frontier interpretable AI models.

Tensors

In this work, a tensor of order $k$ is an array with $k$ indices, $T \in \mathbb{R}^{n_1 \times \dots \times n_k}, n_i,k\in \mathbb{N}$. A scalar has order $0$, a vector order $1$, a matrix order $2$, and so on. The table below lists common tensors, how to generate each at random in PyTorch, and its geometric meaning. In our diagrams, a tensor is a circle with “legs”. Each leg stands for one geometric dimension, that is, one array index.

Name / op	Array shape	Algebraic
Scalar `s = torch.randn(1)`	$\left[\begin{smallmatrix}\cdot\end{smallmatrix}\right]$	$s \in \mathbb{R}$
Vector `v = torch.randn(n)`	$\left[\begin{smallmatrix}\cdot\\\cdot\\\cdot\end{smallmatrix}\right]$	$v \in \mathbb{R}^n$
Matrix `A = torch.randn(m, n)`	$\left[\begin{smallmatrix}\cdot&\cdot\\\cdot&\cdot\\\cdot&\cdot\end{smallmatrix}\right]$	$A \in \mathbb{R}^{m \times n}$
3-Tensor `T = torch.randn(a, b, c)`	—	$T \in \mathbb{R}^{a \times b \times c}$

Operations

AI models manipulate tensors using tensor operations. Tensor operations admit a formal notation known as Einstein-inspired notation for operations or “Einops” . PyTorch supports two main Einops: rearrange, which reorders tensor’s axes, and einsum, which combines multiple tensors. Given their expressivity, mastering Einops, and einsum above all, is one of the most useful skills for designing and implementing AI architectures.

Unfortunately, PyTorch einsum is hard to parse at a glance. It is not the best tool for designing, comparing, or reasoning about tensor operations. It does, however, map one to one onto Penrose graphical notation. This visual notation makes complex tensor operations clear, formal, and unambiguous.

The convention works as follows . To combine two tensors, we connect legs that share a label; a shared label marks the same geometric dimension/index. A connected pair of legs is contracted: we multiply the two tensors’ entries together for each value of that shared index, then sum over the index. This leaves “cancels” the leg out. Any legs left unconnected are “free” legs, and they become the indices of the result.

We can turn each diagram directly into PyTorch code with einsum, using the convention einsum('legs_of_input_tensor_A,legs_of_input_tensor_B->legs_of_output_tensor', A, B). This lets us design a tensor operation as a diagram, get the diagram’s clarity, and then convert it straight into working PyTorch code.

The table below lists the most common tensor operations. These form the building blocks of the tensor manipulations used in frontier interpretable models.

Name / op	Diagram	Algebraic
Scalar multiplication `h = torch.einsum(',->', v, w)`		$h = a b$
Element-wise product `h = einsum('i,i->i', a, b)`		$h_i = a_i b_i$
Dot product `h = einsum('i,i->', a, b)`		$h = \sum_i a_i b_i$
Squared norm `h = einsum('i,i->', a, a)`		$h = \sum_i a_i a_i = \\|a\\|_2$
Normalize `asqn = einsum('i,i->', a, a)` `an = a / asqn**0.5`		$\hat a_i = a_i \big/ \sqrt{\sum_j a_j^2}$
Convex combination `h = einsum('i,i->', a, bn)`		$h = \sum_i a_i \hat b_i$
Cosine similarity `h = einsum('i,i->', an, bn)`		$h = \sum_i \hat a_i \hat b_i$
Sum of matrix slices `h = einsum('ik->i', A)`		$h_i = \sum_k A_{ik}$
Matrix-vector product `h = einsum('ij,i->j', A, b)`		$h_j = \sum_i A_{ij}b_i$
Scaled matrix slices `h = einsum('ij,i->ij', A, b)`		$H_{ij} = A_{ij} b_i$
Matrix-matrix product `H = einsum('ij,ik->jk', A, B)`		$H_{jk} = \sum_i A_{ij}B_{ik}$

Graphical design of simple neural models

With linear algebra fresh in mind, we can now use the graphical notation to design AI models. As introductory examples, we design two familiar models, a linear model and a multi-layer perceptron , before moving to more advanced cases, such as self-attention .

A linear model is one of the oldest models in statistics, yet it remains an important baseline for interpretable machine learning, and it forms the backbone of more complex operations in frontier models. A linear model is a matrix-vector product followed by an activation function. The vector $x \in \mathbb{R}^d$ holds the features of an input sample, and the matrix $W \in \mathbb{R}^{h \times d}$ holds the model’s learnable parameters. Since this model usually has a non-linear activation which makes the diagram asymmetric, we extend Penrose diagrams drawing the input node in gray and performing tensor operations from left to right (or top-down):

Code	Algebraic
`y = sigma(einsum('j,ij->i', x, W))`	$y_i = \sigma \left(\sum_j W_{ij}x_j \right)$

We can apply a linear model to many inputs at once by stacking samples $x_j$ into a batch tensor $X \in \mathbb{R}^{b \times d}$. We can also stack several linear models on top of each other. This gives a multi-layer perceptron (MLP) :

Code	Algebraic
`z1 = sigma(einsum('bj,ij->bi', X, W0))` `z2 = sigma(einsum('bj,ij->bi', z1, W1))` … `y = sigma(einsum('bj,ij->bi', zL, WL))`	$$ Y_{bi_L} = \sigma \left( \sum_{i_{L-1}} W_{i_L, i_{L-1}}^{(L)} \dots \sigma \left( \sum_{j_0} W_{i_0, j_0}^{(0)} X_{bj_0} \right) \right) $$

Self-attention is a key, more complex operation in frontier AI models. This operation projects an input sequence $Z$ of $t$ tokens into a query $q$, key $k$, and value $v$ embeddings. For each pair of tokens $(t,t_p)$, self-attention scores how relevant token $t_p$ is to token $t$. It then uses these relevance scores as weights to combine the value vectors.

Since the tensor manipulations are a bit more complex, we break down the self-attention mechanism into simple atomic manipulations. The first step is to “copy” the input $Z$ since we need to reuse this tensor multiple times. In our notation, copying a tensor can be expressed by branching all its legs. We use the index $t_p$ for the legs of the second and third copy of $Z$ as these legs will be used to index key and value tokens the query can attend to:

Each copy of the tensor $Z$ gets multiplied by a matrix $W \in \mathbb{R}^{d \times e}$ to produce key, query, and value tensors $k,q,v \in \mathbb{R}^{t \times e}$:

For each pair of tokens $(t,t_p)$, we compute how much the query token $t$ attends to the key token $t_p$:

We then normalise these “affinity” scores into probability values using a softmax activation:

And finally we can compute the new embedding of the token $t$ as a convex combination of value embeddings $v$ weighted by their respective probability score:

In a single diagram we can draw self-attention as follows:

Code	Algebraic
`q = einsum('td,de->te', Z, W_q)` `k = einsum('pd,de->pe', Z, W_k)` `v = einsum('pd,de->pe', Z, W_v)`	$q_{te} = \sum_d Z_{td} W_{q,de}$ $k_{pe} = \sum_d Z_{pd} W_{k,de}$ $v_{pe} = \sum_d Z_{pd} W_{v,de}$
`l = einsum('te,pe->tp', q, k) / sqrt(e)`	$l_{tp} = \frac{1}{\sqrt{e}} \sum_e q_{te} k_{pe}$
`probs = softmax(l, dim=-1)`	$\text{probs}_{tp} = \text{softmax}_p \left( l_{tp} \right)$
`h = einsum('tp,pe->te', probs, v)`	$h_{te} = \sum_p \text{probs}_{tp} v_{pe}$

From here on, diagrams stay minimal: we draw only the indices involved in a contraction. PyTorch supports this directly through ellipsis notation, which lets a tensor operation generalize to any number of batch dimensions. For example, an operation with three preserved indices, batch $b$, token $t$, and head $q$, written as

einsum('btqij,btqjk->btqik', A, B)

can be rewritten as

einsum('...ij,...jk->...ik', A, B)

Graphical design of interpretable architectures

Interpretable architectures can be generally segmented into three distinct components : a backbone that maps input $x$ to a hidden representation $z$, a concept encoding map that turns $z$ into human-meaningful concepts $c$, and a concept composition map that turns those concepts into a task prediction $y$.

Most interpretable architectures use specific tensor operations in their concept encoding and concept composition maps to meet interpretability constraints . Here we analyse the most common and recurring of these operations, shared across different families of interpretable models.

Concept encoding maps

Concept encoding maps transform latent representations $z$ into representations $c$, known as concepts, that are constrained to align with human semantics. The most common maps in the literature, in order of increasing tensor-manipulation complexity, are probes such as concept activation vectors (CAVs) and sparse autoencoders (SAEs) , concept bottlenecks , and prototype-based models .

Sparse encoders map latent representations $z \in \mathbb{R}^d$ into the sparse activations $c \in \mathbb{R}^k$ via a sparse linear map $W \in \mathbb{R}^{d \times k}$ with $k \gg d$

Code	Algebraic
`c = sigma(einsum('d,dk->k', z, W))`	$c_k = \sigma \left(\sum_d W_{kd}z_d \right)$

Concept bottlenecks map a latent representation $z$ into the concept representation $c$ via a supervised linear map

Code	Algebraic
`c = sigma(einsum('d,dk->k', z, W))`	$c_k = \sigma \left(\sum_d W_{kd}z_d \right)$

In both cases the tensor operation is identical. The difference lies in the loss and in what the concepts mean: sparse probes recover concept semantics post-hoc (through additional data and labels), while concept bottlenecks build concept semantics into the loss from the start using ground-truth concept annotations $c^{[h]}$.

Concept embedding bottlenecks map a latent representation $z$ into a high dimensional concept representation $u \in \mathbb{R}^{d \times k \times s \times e}$ where $s$ is the concept cardinality and $e$ the embedding size. This concept representation is then used to compute concept predictions $c_k$:

Code	Algebraic
`u = einsum('d,dkse->kse', z, W)`	$u_{kse} = \sum_d W_{dske}z_d$
`c = sigma(einsum('kse,e->ks', u, S))`	$c_{ks} = \sigma \left( \sum_{e} u_{kse} S_{e} \right)$

Prototype-based concept maps need a genuinely different tensor operation. We can think of prototypes as reference examples that tell us whether a concept is active. For instance, the embedding of an apple or a ball can serve as a positive “prototypical example” for the concept round, and a fridge or a book as a negative example. Ground-truth prototype labels sit in the tensor $\pi^{[h]} \in \mathbb{R}^{p \times k}$, so each concept $k$ has $p$ labelled prototypes. For a concept $k$ and an input embedding $z \in \mathbb{R}^d$, we compare $z$ against every prototype in $P \in \mathbb{R}^{d \times p \times k}$ and compute the concept label based on input-prototype similarity. For instance, if $z$ is closer to the prototypes for book and fridge than to the prototypes for apple and ball, then the predicted label for round should sit close to $0$.

Code	Algebraic
`l = einsum('d,dpk->pk', zn, Pn)`	$\text{l}_{pk} = \sum_d \hat{P}_{dpk}z_d$
`probs = softmax(l, dim=-1)`	$\text{probs}_{pk} = \text{sm} \left(l_{pk} \right)$
`c_pred = einsum('pk,pk->k', probs, p_true))`	$c_k = \sum_p \text{probs}_{pk}\pi_{pk}^{[h]}$

Concept composition maps

In most interpretability works, the concept composition map is a simple linear model: self-explaining neural nets , sparse autoencoders , concept bottleneck models , all use linear models. A few exceptions are worth discussing: neural additive models , concept embedding predictors , and mixtures of linear models .

Neural additive models transform concept activations $c$ independently using a different MLP for each concept $k$ and output task $r$. Then, for each task, they sum the outputs of the MLP of each concept to predict target $y_r$:

Code	Algebraic
`h1 = sigma(einsum('k,kri->kri', c, W1))`	$h_{kri}^{(1)} = \sigma \left( c_k W^{(1)}_{kri} \right)$
`h2 = sigma(einsum('kri,krij->krj', h1, W2))`	$h_{krj}^{(2)} = \sigma \left( \sum_i h_{kri}^{(1)} W^{(2)}_{krij} \right)$
…	…
`hL = einsum('kre,kre->kr', he, WL)`	$h_{kr}^{(L)} = \sum_e h_{kre}^{(L-1)} W^{(L)}_{kre}$
`y = sigma(einsum('kr->r', hL))`	$y_r = \sigma \left(\sum_k h_{kr}^{(L)}\right)$

Concept embedding predictors rescale concept embeddings $u$ (e.g., generated by a concept embedding bottleneck) using concept activations $c$ before projecting into the output space $r$:

Code	Algebraic
`h = einsum('kse,ks->kse', u, c)`	$h_{kse} = u_{kse} c_{ks}$
`h = einsum('kse->ke', h)`	$h_{ke} = \sum_s h_{kse}$
`l = einsum('ke,ker->r', h, W)`	$l_r = \sum_{ke} h_{ke} W_{ker}$
`l = sigma(l)`	$y_r = \sigma \left(l_r \right)$

Mixtures of linear models compute different predictions for the target $y_r$ using $m$ different linear models. Then each prediction is weighted by the probability of selecting a specific linear model:

Code	Algebraic
`l = einsum('d,dm->m', z, W))`	$l_m = \sum_d W_{dm} z_d$
`pr = softmax(l, dim=-1)`	$\text{pr}_m = \text{sm} \left( l_m \right)$
`v = einsum('k,kmr->mr', c, E)`	$v_{mr} = \sum_k E_{kmr} c_k$
`ly = einsum('m,mr->r', pr, v)`	$l_r = \sum_{m} v_{mr} \text{pr}_m$
`y = sigma(ly)`	$y_r = \sigma \left(l_r \right)$

Case study: frontier interpretable language models

As a case study, we diagram the architecture of Steerling-8B , the largest interpretable-by-design language model publicly available at the time of writing.

To keep the focus on the essential tensor manipulations, we drop batch dimensions, since they aren’t involved in any contraction, and we show a single attention head; extending to multiple heads is straightforward. Under these conditions, the essential tensor manipulations in the Steerling-8B architecture take about 30 lines of code. Drawing the Steerling-8B tensor diagram has three main benefits over the notation used in the original technical report :

It shows at a glance that Steerling-8B is a residual model: the gradient can flow from the output straight back to the first input.
It maps one to one onto PyTorch einops and activation functions, which makes the diagram useful for reproducing the model faithfully.
Every operation in the diagram has a direct geometric reading in linear algebra. This lets us read each manipulation as a transformation in space, which helps build intuition for the underlying computation.

ht = einsum('tr,re->te', X, E)
hp = einsum('ti,ie->te', X, Ep)

Z = ht + hp

q = einsum('te,ed->td', Z, W_q)
k = einsum('pe,ed->pd', Z, W_k)
v = einsum('pe,ed->pd', Z, W_v)

l = einsum('te,pe->tp', q, k)
l = l / sqrt(W_k.shape[1])

l = l + M

probs = softmax(l, dim=-1)

h = einsum('tp,pe->te', probs, v)

h = einsum('te,ed->td', h, W)
h = dropout(h)

Z = Z + h

h = layernorm(Z)

h = einsum('te,ed->td', h, W1)

h = sigma(h)

h = einsum('te,ed->td', h, W2)

h = dropout(h)

Z = Z + h

ls = einsum('te,ek->tk', Z, Ws)
lu = einsum('te,eu->tu', Z.detach(), Wu)

cs = sigmoid(ls)
cu = sigmoid(lu)

csf = topk(cs)
cuf = topk(cu)

csfe = einsum('tk,ke->te', csf, Ks)
cufe = einsum('tu,ue->te', cuf, Ku)

ce = csfe + cufe

Z = Z - ce

Z = Z + ce

l = einsum('te,er->tr', Z, Wh)

y = sigma(l)

Discussion

Graphical tensor notation dates to Penrose , who introduced diagrams for tensor contraction in physics. The notation has since been adopted by the categorical-quantum-mechanics community , and, more recently, by theoretical computer science and machine learning.

A first line of work has proposed general-purpose diagrammatic languages for deep learning architectures, without a focus on interpretability. Chiang et al. propose named-axis tensor notation to disambiguate operations such as attention. Abbott introduces neural circuit diagrams, a graphical language with a formal correspondence to implementation, later used to derive memory-efficient attention algorithms . Cruttwell et al. pursue a category-theoretic account of architectures more broadly, using string diagrams, a mathematical generalization of Penrose notation, to unify architectures such as convolutional neural nets, recurrent neural nets, and transformers under one algebraic framework.

A more recent line of work started analysing the interpretability literature using graphical notations. Giannini et al. , Tull et al. , and Barbiero et al. use string diagrams to analyse explainable AI methods and interpretable architectures, but without using the tensor manipulation semantics that maps directly to PyTorch programming interfaces. Taylor applies Penrose notation to mechanistic interpretability, using it to reverse-engineer trained transformer components such as induction heads, building on the informal flowcharts used by Elhage et al. to describe transformer circuits. However, this line of work analyses pre-trained opaque models and does not consider tensor manipulations required by inherently interpretable models.

This paper takes a notation developed for post-hoc analysis of trained models and, for the first time to our knowledge, applies this notation to the forward problem: analysing and designing architectures that are interpretable by construction, with a direct, mechanical path from diagram to PyTorch code.

Limitations and concrete usage

Tensor diagrams are exact for multilinear operations, but nonlinearities, masking, and discrete operations such as top-k require ad hoc extensions. Closeness to implementation is also a double-edged sword: diagram size grows with the complexity of the tensor manipulations, so full frontier architectures quickly become unwieldy to draw in conference papers. For this reason, we see tensor diagrams as best paired with a coarser formalism such as probabilistic graphical models. Probabilistic graphical models may be used to capture the high-level causal structure between random variables, while small tensor diagrams specify how each conditional probability function is implemented.

Conclusion

We have shown how tensor diagrams can guide the design of interpretable deep neural networks. For the most common tensor manipulations in interpretability research, we have built a “Rosetta stone” showing the matching diagram, PyTorch code, geometric interpretation, and symbolic equation side by side, so readers from different backgrounds can compare and understand them. We also tackled a harder case: we diagrammed and implemented the key modules of a frontier interpretable-by-design language model, Steerling-8B in about 30 lines of code.

Tensor diagrams are expressive, formal, and map directly onto PyTorch code. For these reasons, they could become a standard tool for designing, comparing, and implementing interpretable architectures, alongside other graphical tools such as probabilistic graphical models.

Name / op	Array shape	Algebraic
Scalar `s = torch.randn(1)`	\(\left[\begin{smallmatrix}\cdot\end{smallmatrix}\right]\)	\(s \in \mathbb{R}\)
Vector `v = torch.randn(n)`	\(\left[\begin{smallmatrix}\cdot\\\cdot\\\cdot\end{smallmatrix}\right]\)	\(v \in \mathbb{R}^n\)
Matrix `A = torch.randn(m, n)`	\(\left[\begin{smallmatrix}\cdot&\cdot\\\cdot&\cdot\\\cdot&\cdot\end{smallmatrix}\right]\)	\(A \in \mathbb{R}^{m \times n}\)
3-Tensor `T = torch.randn(a, b, c)`	—	\(T \in \mathbb{R}^{a \times b \times c}\)

Name / op	Diagram	Algebraic
Scalar multiplication `h = torch.einsum(',->', v, w)`		\(h = a b\)
Element-wise product `h = einsum('i,i->i', a, b)`		\(h_i = a_i b_i\)
Dot product `h = einsum('i,i->', a, b)`		\(h = \sum_i a_i b_i\)
Squared norm `h = einsum('i,i->', a, a)`		\(h = \sum_i a_i a_i = \\|a\\|_2\)
Normalize `asqn = einsum('i,i->', a, a)` `an = a / asqn**0.5`		\(\hat a_i = a_i \big/ \sqrt{\sum_j a_j^2}\)
Convex combination `h = einsum('i,i->', a, bn)`		\(h = \sum_i a_i \hat b_i\)
Cosine similarity `h = einsum('i,i->', an, bn)`		\(h = \sum_i \hat a_i \hat b_i\)
Sum of matrix slices `h = einsum('ik->i', A)`		\(h_i = \sum_k A_{ik}\)
Matrix-vector product `h = einsum('ij,i->j', A, b)`		\(h_j = \sum_i A_{ij}b_i\)
Scaled matrix slices `h = einsum('ij,i->ij', A, b)`		\(H_{ij} = A_{ij} b_i\)
Matrix-matrix product `H = einsum('ij,ik->jk', A, B)`		\(H_{jk} = \sum_i A_{ij}B_{ik}\)

Code	Algebraic
`q = einsum('td,de->te', Z, W_q)` `k = einsum('pd,de->pe', Z, W_k)` `v = einsum('pd,de->pe', Z, W_v)`	\(q_{te} = \sum_d Z_{td} W_{q,de}\) \(k_{pe} = \sum_d Z_{pd} W_{k,de}\) \(v_{pe} = \sum_d Z_{pd} W_{v,de}\)
`l = einsum('te,pe->tp', q, k) / sqrt(e)`	\(l_{tp} = \frac{1}{\sqrt{e}} \sum_e q_{te} k_{pe}\)
`probs = softmax(l, dim=-1)`	\(\text{probs}_{tp} = \text{softmax}_p \left( l_{tp} \right)\)
`h = einsum('tp,pe->te', probs, v)`	\(h_{te} = \sum_p \text{probs}_{tp} v_{pe}\)

Code	Algebraic
`u = einsum('d,dkse->kse', z, W)`	\(u_{kse} = \sum_d W_{dske}z_d\)
`c = sigma(einsum('kse,e->ks', u, S))`	\(c_{ks} = \sigma \left( \sum_{e} u_{kse} S_{e} \right)\)

Code	Algebraic
`l = einsum('d,dpk->pk', zn, Pn)`	\(\text{l}_{pk} = \sum_d \hat{P}_{dpk}z_d\)
`probs = softmax(l, dim=-1)`	\(\text{probs}_{pk} = \text{sm} \left(l_{pk} \right)\)
`c_pred = einsum('pk,pk->k', probs, p_true))`	\(c_k = \sum_p \text{probs}_{pk}\pi_{pk}^{[h]}\)

Code	Algebraic
`h1 = sigma(einsum('k,kri->kri', c, W1))`	\(h_{kri}^{(1)} = \sigma \left( c_k W^{(1)}_{kri} \right)\)
`h2 = sigma(einsum('kri,krij->krj', h1, W2))`	\(h_{krj}^{(2)} = \sigma \left( \sum_i h_{kri}^{(1)} W^{(2)}_{krij} \right)\)
…	…
`hL = einsum('kre,kre->kr', he, WL)`	\(h_{kr}^{(L)} = \sum_e h_{kre}^{(L-1)} W^{(L)}_{kre}\)
`y = sigma(einsum('kr->r', hL))`	\(y_r = \sigma \left(\sum_k h_{kr}^{(L)}\right)\)

Code	Algebraic
`h = einsum('kse,ks->kse', u, c)`	\(h_{kse} = u_{kse} c_{ks}\)
`h = einsum('kse->ke', h)`	\(h_{ke} = \sum_s h_{kse}\)
`l = einsum('ke,ker->r', h, W)`	\(l_r = \sum_{ke} h_{ke} W_{ker}\)
`l = sigma(l)`	\(y_r = \sigma \left(l_r \right)\)

Code	Algebraic
`l = einsum('d,dm->m', z, W))`	\(l_m = \sum_d W_{dm} z_d\)
`pr = softmax(l, dim=-1)`	\(\text{pr}_m = \text{sm} \left( l_m \right)\)
`v = einsum('k,kmr->mr', c, E)`	\(v_{mr} = \sum_k E_{kmr} c_k\)
`ly = einsum('m,mr->r', pr, v)`	\(l_r = \sum_{m} v_{mr} \text{pr}_m\)
`y = sigma(ly)`	\(y_r = \sigma \left(l_r \right)\)