Hi! I’m experimenting with a project
Einstoff
;
it is inspired by einops/einx, but instead of parsing a string with syntax,
it uses native WL patterns as a eDSL to describe tensor operations like
reshaping, reducing, contracting, joining, and splitting multidimensional arrays
Here are some basic examples:
In[]:=
ims=MapIndexed[ImageMultiply[{#1,​​Lighter[ColorData["DefaultChartColors"]@First@#2,.7]}]​​&,Rasterize[#,RasterSize->56]&/@Characters@"Einstoff"]​​ImageData/@ims//Dimensions
Out[]=

,
,
,
,
,
,
,

Out[]=
{8,120,56,3}
In[]:=
[◼]
GitHubInstall
["Gravifer","Einstoff.wl","v0.1.0-alpha.1"]​​Needs["Einstoff`"]
Out[]=
PacletObject
Name: Einstoff
Version: 0.1.0

Einstoff[ArrayReshape]
mostly does what
Transpose
and
ArrayReshape
do; you specify the shapes of the input and output tensors (mostly) symbolically,
In[]:=
img=Einstoff[ArrayReshape][{{batch_,height_,width_,color_}}:>{{height,batch⊗width,color}},{ImageData/@ims}];​​img//Image//Thumbnail
Out[]=
Under the hood, it actually does
In[]:=
Einstoff[ArrayReshape]{{batch_,height_,width_,color_}}:>{{height,batch⊗width,color}},
NumericArray[
]
,TraceAction->Defer
Out[]=
ArrayReshape
T:{2,1,3,4}
NumericArray
Type: Real64
Dimensions: {8,120,56,3}
Data not saved. Save now

,{120,448,3}
You can also massage the tensor in some other commonly encountered ways:
In[]:=
Einstoff[ArrayReduce][Mean][{{h_,"b"⊗w_,c_}}:>{{h,w,c}},{img},{"b"->8}]//Image//Thumbnail
Out[]=
In[]:=
Einstoff[ArrayReshape][{{b1⊗b2,h1_⊗dh,w_,c_}}:>{{h1,(b2⊗b1)⊗(dh⊗w),c}},{ImageData/@ims},{b1->4,dh->2}]//Image//Thumbnail
Out[]=
In[]:=
Einstoff[Operate][Reverse]
h_
,"b"⊗w_,c_:>{{h,"b"⊗w,c}},{img},{"b"->8}//Image//Thumbnail
Out[]=
In[]:=
Thumbnail@*Image/@Einstoff[Split][{{height_,width_,"r"⊕"g"⊕"b"}}:>{{height,width},{height,width},{height,width}},{ImageData@ExampleData[{"TestImage","Apples"}]}]
Out[]=

,
,

When the building blocks are simple but the routine is hairy, this interface can read more concise than explicit procedural code.
Does it feel intuitive or idiomatic? Is the programming interface awkward?
The project is still in extremely early stage, and I’m struggling to get the interface consistent and the documentation synced.
I think it won’t take long to run into rough edges when trying it out, but I will keep trying to improve it!

Motivation

A lot of image processing and machine learning techniques rely on structured manipulations of multi-dimensional arrays, or loosely, tensors.
Being a list-native language, Wolfram Language has a
plethora
of convenient and performant primitives to handle multi-dimensional arrays;and, even better, these built-in functions almost always work smoothly with any representation of arrays,no matter if they are
nested-lists
with all M-expression flexibility,
symbolic tensors
, declarative and efficient
sparse arrays
, or
packed
objects
that have native speed.
With batching, striding, resampling and so on, array data can often contain dimensions that are mapped over rather than operated upon, and these can often need joining, splitting or recombining on the fly; even when the individual operations are simple, the imperative flow of code can often grow hard to reason about, requiring verbose variable names, comments, and exterior documentation.

Einstein notation

For linear algebraic operations, the Einstein summation convention is a great convenience; mathematicians and physicians enjoy it so much that they extended it into Ricci calculus and abstract index notation. Programmers tapped into it as well:
In[]:=
a=ArrayReshape[Range[6],{3,2}];​​b=ArrayReshape[Range[12],{4,3}];​​(b.a)
Out[]=
{{22,49,76,103},{28,64,100,136}}
〉
In[]:=
import numpy as np
​
a = np.array(<* a *>)
b = np.array(<* b *>)
np.einsum('ki,jk->ij', a, b).tolist()
Out[]=
{{22,49,76,103},{28,64,100,136}}
〉
In[]:=
ra = np.random.randn(8,120,56,3)
rb = np.random.randn(8,120,56,3)
np.einsum('b ..., b ... -> ...', ra, rb) # dot product, batched with an anonymous sequence of axes
Out[]=
NumericArray
Type: Real64
Dimensions: {120,56,3}

Which the
[◼]
EinsteinSummation
ResourceFunction made even better, using a
Rule
instead of a string for the descriptive:
In[]:=
[◼]
EinsteinSummation
[{{k,i},{j,k}}->{i,j},{a,b}]
Out[]=
{{22,49,76,103},{28,64,100,136}}
The resource function allows using any distinct expression for axis indices;
by not parsing a string and building an AST oneself, the DSL can be easily extended due to the inherent expressiveness of WL.
In[]:=
SeedRandom[1];​​u=RandomReal[1,{3,3}];​​x=RandomReal[1,{3,3}];​​
[◼]
EinsteinSummation
[{{1,2},{3,2}},{u,x}]
Out[]=
{{1.18725,1.14471,1.23396},{0.231893,0.203386,0.416294},{0.725267,0.673465,0.890237}}
In[]:=
[◼]
EinsteinSummation
[{{i,k},{j,k}}->{j,i},{A,M}]
Out[]=
TensorTranspose[TensorContract[AM,{{2,4}}],{2,1}]

Einsnein notation

Having the axes of multi-dimensional arrays named is expressive and self-documenting,
so people have tried to generalize it to operations other than
Inner
contractions, like reshaping or reduction.
Enter einops:
〉
In[]:=
from einops import rearrange​ims = np.array(<* ImageData /@ ims *>)​rearrange(ims[0], "h w c -> w h c") # equivalent to np.transpose(ims[0], (1, 0, 2))

In[]:=
{%//Image//Thumbnail,%==Transpose[ImageData@ims[[1]]]}
Out[]=

,True
And it’s also possible to use more verbose axis names, leading to well self-documenting code:
〉
In[]:=
rearrange(ims, 'batch height width color -> height (batch width) color') # the `(batch width)` part means combining these into one axis, as a direct product

In[]:=
%//Image//Thumbnail
Out[]=
〉
In[]:=
from einops import reduce​# reducing an axis by taking the meanreduce(ims, "b h w c -> h w c", "mean") # `b` is dropped from RHS, so it is averaged over; equivalent to np.mean(x, -1)

Python and WL share the row-major convention when it comes to arrays, i.e., the later an axis comes, the more rapidly its index varies.
This means that the order of a parenthesized dimension matters:
Steering into it wisely means flexible striding and leads to an impressive gallery.
◼
  • Immediate literal axes
  • 〉
    In[]:=
    import einx
    ​
    y = einx.id("a b -> a b 3", x)
    # same as
    y = einx.id("a b -> a b c", x, c=3)
    # or
    from einops import repeat
    repeat(x, "a b -> a b c", c=3)
    ◼
  • Explicitly targeted axes for operation (compare with the NumPy example above)
  • 〉
    In[]:=
    np.einsum("ki,jk->ik", a, b).tolist() # a typo fails silently
    # It computes sum-reduction along j
    # followed by element-wise multiplication.
    With einx
    〉
    In[]:=
    einx.dot("[k] i, j [k] -> i j", a, b).tolist() # correct spelling
    typos are checked and become exceptions
    〉
    In[]:=
    einx.dot("[k] i, j [k] -> i k", a, b).tolist()
    〉
    In[]:=
    einx.dot("k i, j k -> i k", a, b).tolist()
    ◼
  • Named binding to a sequence of axes
  • 〉
    In[]:=
    x = np.random.randn(10, 20, 30, 40)
    einx.sum("s... [c] -> s...", x) # reduction by total; handled as einx.sum("s1 s2 s3 [c] -> s1 s2 s3", x)
    〉
    In[]:=
    x = np.random.randn(10)
    einx.sum("s... [c] -> s...", x) # the same operation works, with variadic expressibility
    〉
    In[]:=
    x = np.random.randn(10, 20)
    y = np.random.randn(10, 20)
    ​
    einx.add("s..., s... -> s...", x, y) # naming the sequence allows repeated matching
    〉
    In[]:=
    # by naming ellipses, they can match with different sequences
    ​
    x = np.random.randn(10, 20)
    y = np.random.randn(30, 40)
    ​
    einx.add("a..., b... -> a... b...", x, y) # equivalent to a batched `Outer[Plus,...]`
    What’s more, the binding can be structural:
    〉
    In[]:=
    x = np.random.randn(640, 480, 3)
    ​
    # Divide an image into a (flattened) list of 4x4 patches
    ​
    y = einx.id("(s ds)... c -> (s...) ds... c", x, ds=4)
    # expands to
    # einx.id("(s1 ds1) (s2 ds2) c -> (s1 s2) ds1 ds2 c", x, ds1=4, ds2=4)
    y
    〉
    In[]:=
    x = np.random.randn(640, 480)
    ​
    # Perform mean pooling with 4x4 patches (if evenly divisible)
    ​
    y = einx.mean("(s [ds])...", x, ds=4)
    # expands to
    # einx.mean("(s1 [ds1]) (s2 [ds2]) -> s1 s2", x, ds1=4, ds2=4)
    y

    M-expr is a good AST

    Current status

    The project at this point is just a barely-working POC, and is grossly lacking compared to einx/einops or existing WL resource functions, feature-wise, testing-wise, etc.
    Some essential features that are entirely absent right now:
    ◼
  • Indexing and gather/scattering of entries;
  • ◼
  • Bit-mask / gate controlled operation;
  • ◼
  • Optimized lowered computation graph;
  • ◼
  • Within-tensor contraction (generalized traces) and diagonal extraction
  • and many stuff that are more obvious but still not yet implemented.
    Though, the goal is not to clone einops or einx feature-for-feature either;
    I am more interested in whether tensor-axis notation can be made to feel idiomatic in Wolfram Language:
    using pattern bindings for inferred dimensions, symbolic expressions as an AST, and the existing array primitives as the lowering targets.
    I hope this makes more able people find this interesting or worth looking into.
    I’m especially interested in feedback from experienced Wolfram users on the interface:
    ◼
  • Does this feel natural, or too foreign despite being made of native WL expressions?
  • I would be grateful for API/design feedback before trying to polish this further for wider release,
    and any thoughts or comments are much obliged ☺