Ecosyste.ms: Awesome
An open API service indexing awesome lists of open source software.
https://github.com/nicholas-miklaucic/eins
One tensor operation is all you need
https://github.com/nicholas-miklaucic/eins
array deep-learning einops jax numpy python
Last synced: 6 days ago
JSON representation
One tensor operation is all you need
- Host: GitHub
- URL: https://github.com/nicholas-miklaucic/eins
- Owner: nicholas-miklaucic
- License: mit
- Created: 2024-03-14T01:20:59.000Z (7 months ago)
- Default Branch: main
- Last Pushed: 2024-04-24T01:29:44.000Z (5 months ago)
- Last Synced: 2024-07-10T15:25:13.506Z (3 months ago)
- Topics: array, deep-learning, einops, jax, numpy, python
- Language: Python
- Homepage: https://nicholas-miklaucic.github.io/eins/
- Size: 2.9 MB
- Stars: 2
- Watchers: 1
- Forks: 0
- Open Issues: 4
-
Metadata Files:
- Readme: README.md
- Contributing: CONTRIBUTING.md
- License: LICENSE.txt
- Code of conduct: CODE_OF_CONDUCT.md
Awesome Lists containing this project
README
# eins
## One tensor operation is all you need
[](https://pypi.org/project/eins)
[](https://pypi.org/project/eins)
[](code_of_conduct.md)
[](https://makeapullrequest.com)
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:
```python
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](https://nn.labml.ai/transformers/vit/index.html#PatchEmbeddings). That means breaking
up an image into patches and then linearly embedding each patch.
```python
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`](https://github.com/arogozhnikov/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](https://nicholas-miklaucic.github.io/eins/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](https://nicholas-miklaucic.github.io/eins/) or
the [examples](examples/README.md).
## Installation
```console
pip install eins
```
Eins works with anything that implements the [Array
API](https://data-apis.org/array-api/latest/index.html), 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`
- [x] 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`](https://github.com/arogozhnikov/einops) library
inspired this project and its syntax. After working on my own extension to
handle indexing, I realized that
[`eindex`](https://github.com/arogozhnikov/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](https://spdx.org/licenses/MIT.html) license.