Package: einx
einx- Tensor Operations in Einstein-Inspired Notationeinx is a Python library that allows formulating many tensor operations as concise expressions using Einstein notation. It is inspired byeinops, but follows a novel and unique design:Fully composable and powerful Einstein expressions with[]-notation.Support for many tensor operations (einx.{sum|max|where|add|dot|flip|get_at|...}) with Numpy-like naming.Easy integration and mixing with existing code. Supports tensor frameworks Numpy, PyTorch, Tensorflow and Jax.Just-in-time compilation of all operations into regular Python functions using Python'sexec().Optional:Generalized neural network layers in Einstein notation. Supports PyTorch, Flax, Haiku, Equinox and Keras.Getting started:TutorialExample: GPT-2/ Mamba with einxHow does einx compare with einops?API referenceInstallationpipinstalleinxSeeInstallationfor more information.What does einx look like?Tensor manipulationimporteinxx={np.asarray|torch.as_tensor|jnp.asarray|tf.convert_to_tensor}(...)# Create some tensoreinx.sum("a [b]",x)# Sum-reduction along columnseinx.flip("... (g [c])",x,c=2)# Flip pairs of values along the last axiseinx.mean("b [s...] c",x)# Global mean-poolingeinx.sum("b (s [s2])... c",x,s2=2)# Sum-pooling with kernel_size=stride=2einx.add("b... [c]",x,b)# Add biaseinx.get_at("b [h w] c, b i [2] -> b i c",x,y)# Gather values at coordinateseinx.rearrange("b (q + k) -> b q, b k",x,q=2)# Spliteinx.rearrange("b c, 1 -> b (c + 1)",x,[42])# Append number to each channeleinx.dot("... [c1|c2]",x,y)# Matmul = linear map from c1 to c2 channels# Vectorizing mapeinx.vmap("b [s...] c -> b c",x,op=np.mean)# Global mean-poolingeinx.vmap("a [b], [b] c -> a c",x,y,op=np.dot)# MatmulAll einx functions simply forward computation to the respective backend, e.g. by internally callingnp.reshape,np.transpose,np.sumwith the appropriate arguments.Common neural network operations# Layer normalizationmean=einx.mean("b... [c]",x,keepdims=True)var=einx.var("b... [c]",x,keepdims=True)x=(x-mean)*torch.rsqrt(var+epsilon)# Prepend class tokeneinx.rearrange("b s... c, c -> b (1 + (s...)) c",x,cls_token)# Multi-head attentionattn=einx.dot("b q (h c), b k (h c) -> b q k h",q,k,h=8)attn=einx.softmax("b q [k] h",attn)x=einx.dot("b q k h, b k (h c) -> b q (h c)",attn,v)# Matmul in linear layerseinx.dot("b...      [c1|c2]",x,w)# - Regulareinx.dot("b...   (g [c1|c2])",x,w)# - Grouped: Same weights per groupeinx.dot("b... ([g c1|g c2])",x,w)# - Grouped: Different weights per groupeinx.dot("b  [s...|s2]  c",x,w)# - Spatial mixing as in MLP-mixerSeeCommon neural network opsfor more examples.Deep learning modulesimporteinx.nn.{torch|flax|haiku|equinox|keras}aseinnbatchnorm=einn.Norm("[b...] c",decay_rate=0.9)layernorm=einn.Norm("b... [c]")# as used in transformersinstancenorm=einn.Norm("b [s...] c")groupnorm=einn.Norm("b [s...] (g [c])",g=8)rmsnorm=einn.Norm("b... [c]",mean=False,bias=False)channel_mix=einn.Linear("b... [c1|c2]",c2=64)spatial_mix1=einn.Linear("b [s...|s2] c",s2=64)spatial_mix2=einn.Linear("b [s2|s...] c",s=(64,64))patch_embed=einn.Linear("b (s [s2|])... [c1|c2]",s2=4,c2=64)dropout=einn.Dropout("[...]",drop_rate=0.2)spatial_dropout=einn.Dropout("[b] ... [c]",drop_rate=0.2)droppath=einn.Dropout("[b] ...",drop_rate=0.2)Seeexamples/train_{torch|flax|haiku|equinox|keras}.pyfor example trainings on CIFAR10,GPT-2andMambafor working example implementations of language models using einx, andTutorial: Neural networksfor more details.Just-in-time compilationeinx traces the required backend operations for a given call into graph representation and just-in-time compiles them into a regular Python function using Python'sexec(). This reduces overhead to a single cache lookup and allows inspecting the generated function. For example:>>>x=np.zeros((3,10,10))>>>graph=einx.sum("... (g [c])",x,g=2,graph=True)>>>print(graph)# backend: einx.backend.numpydefop0(i0):x1=backend.reshape(i0,(3,10,2,5))x0=backend.sum(x1,axis=3)returnx0SeeJust-in-time compilationfor more details.