GitLocker: The Coding Marketplace

Description:

eins 0.1.1

eins
One tensor operation is all you need

What if most of your machine learning model code could be replaced by a single operation? Eins gives
you a powerful language to describe array manipulation, making it a one-stop shop for all of your AI
needs.
Let's say you want to compute batched pairwise distance between two batches of arrays of points:
from eins import EinsOp
EinsOp('batch n1 d, batch n2 d -> batch n1 n2',
combine='add', reduce='l2_norm')(pts1, -pts2)

If you're more interested in deep learning, here is the patch embedding from a vision
transformer. That means breaking
up an image into patches and then linearly embedding each patch.
from eins import EinsOp
patchify = EinsOp('''b (n_p patch) (n_p patch) c, (patch patch c) emb -> b (n_p n_p) emb''')
patches = patchify(images, kernel)

You input the shapes and Eins will figure out what to do.
If you've used einops, then think of Eins as einops
with a more ambitious goal—being the only operation you should need for your next deep learning
project.
Interested? Check out the tutorial, which
walks you through the highlights of Eins with examples of how Eins can make the array operations you
know and love more readable, portable, and robust.
To learn more, consult the documentation or
the examples.
Installation
pip install eins

Eins works with anything that implements the Array
API, and Eins explicitly promises to support
NumPy, PyTorch, and JAX—including full differentiability. You will need one of those libraries to
actually use Eins operations.
Features

🧩 A solver that can handle duplicate axes, named constants, and constraints
🚀 Compilation and optimization for high performance without sacrificing readability
Split, concatenate, stack, flatten, transpose, normalize, reduce, broadcast, and more
Works across frameworks
A composable set of unified array operations for portable softmax, power normalization, and more

Roadmap
Eins is still in heavy development. Here's a sense of where we're headed.
Near-Term (weeks)

Updating indexing syntax to match eindex
Unit array to indicate zero-dimensional tensors
... for batching over dynamic numbers of batch axes
Specifying intermediate results to control the order of reduction
Support - and /
Better error reporting
Ways of visualizing and inspecting the computation graph
Typo checking in errors about axes
Multiple outputs, either through arrows or commas

Long-Term (months)

Layers for popular ML frameworks
Automatically optimizing the execution of a specific EinsOp for a specific
computer and input size
Completing full support for tensor indexing
Static typing support
Tabulating the model FLOPs/memory usage as a function of named axes
Functionality akin to pack and unpack

Acknowledgements
The excellent einops library
inspired this project and its syntax. After working on my own extension to
handle indexing, I realized that
eindex already had a more coherent
vision for what indexing can look like, and so much of that syntax in this
library borrows from that one.
Contributing
Any contributions to Eins are welcomed and appreciated! If you're interested
in making serious changes or extensions to the syntax of operations, consider
reaching out first to make sure we're on the same page. For any code changes, do
make sure you're using the project Ruff settings.
License
Eins is distributed under the terms of the
MIT license.