Metadata-Version: 2.4
Name: einx
Version: 0.4.0
Summary: Universal Notation for Tensor Operations in Python
Author-email: Florian Fervers <florian.fervers@gmail.com>
License: MIT License
        
        Copyright (c) 2023- Florian Fervers
        
        Permission is hereby granted, free of charge, to any person obtaining a copy
        of this software and associated documentation files (the "Software"), to deal
        in the Software without restriction, including without limitation the rights
        to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
        copies of the Software, and to permit persons to whom the Software is
        furnished to do so, subject to the following conditions:
        
        The above copyright notice and this permission notice shall be included in all
        copies or substantial portions of the Software.
        
        THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
        IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
        FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
        AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
        LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
        OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
        SOFTWARE.
        
Project-URL: Homepage, https://github.com/fferflo/einx
Project-URL: Source, https://github.com/fferflo/einx
Project-URL: Documentation, https://einx.readthedocs.io
Classifier: Development Status :: 3 - Alpha
Classifier: Intended Audience :: Developers
Classifier: License :: OSI Approved :: MIT License
Classifier: Programming Language :: Python :: 3
Requires-Python: >=3.10
Description-Content-Type: text/markdown
License-File: LICENSE
Requires-Dist: numpy
Requires-Dist: sympy
Requires-Dist: frozendict
Provides-Extra: torch
Requires-Dist: torch>=2.2; extra == "torch"
Provides-Extra: mlx
Requires-Dist: mlx; extra == "mlx"
Dynamic: license-file

# *einx* - Universal Notation for Tensor Operations

[![pytest](https://github.com/fferflo/einx/actions/workflows/run_pytest.yml/badge.svg)](https://github.com/fferflo/einx/actions/workflows/run_pytest.yml)
[![Documentation](https://img.shields.io/badge/documentation-link-blue.svg)](https://einx.readthedocs.io)
[![PyPI version](https://badge.fury.io/py/einx.svg)](https://badge.fury.io/py/einx)
[![Python 3.10+](https://img.shields.io/badge/python-3.10+-blue.svg)](https://www.python.org/downloads/release/python-3100/)

einx is a notation and Python library that provides a universal interface to formulate tensor operations in frameworks such as Numpy, PyTorch, Jax, Tensorflow, and MLX.

* [Quickstart](#quickstart)
* [What does einx look like?](#what-does-einx-look-like)
* [How does the notation work?](#how-does-the-notation-work)
* [How are einx operations implemented?](#how-are-einx-operations-implemented)
* [Which operations are supported?](#which-operations-are-supported)

## Quickstart

**Installation:**

```python
pip install einx
```

**Example code:**

```python
import einx
import numpy as np

x = np.ones((10, 20, 30)) # Create some tensor (numpy/torch/jax/tensorflow/mlx/...)

y = einx.sum("a [b] c", x) # Call an einx operation

print(y.shape)
```

**Documentation and tutorials:** [https://einx.readthedocs.io](https://einx.readthedocs.io)

## What does einx look like?

```python
z = einx.id("a (b c) -> (b a) c", x, b=2)             # Permute and (un)flatten axes
z = einx.sum("a [b]", x)                              # Sum-reduction along second axis
z = einx.flip("... (g [c])", x, c=2)                  # Flip pairs of values along the last axis
z = einx.mean("b [...] c", x)                         # Spatial mean-pooling
z = einx.multiply("a..., b... -> (a b)...", x, y)     # Kronecker product
z = einx.sum("b (s [ds])... c", x, ds=(2, 2))         # Sum-pooling with 2x2 kernel
z = einx.add("a, b -> a b", x, y)                     # Outer sum
z = einx.dot("a [b], [b] c -> a c", x, y)             # Matrix multiplication
z = einx.get_at("b [h w] c, b i [2] -> b i c", x, y)  # Gather values at coordinates
z = einx.id("b (q + k) -> b q, b k", x, q=2)          # Split
z = einx.id("b c, -> b (c + 1)", x, 42)               # Append number to each channel
```

See the [documentation](https://einx.readthedocs.io/en/latest/more/operations.html) for more examples.

## How does the notation work?

An einx operation consists of (1) an elementary operation and (2) an einx expression that describes how the elementary operation is vectorized. For example,
the code

```python
z = einx.{OP}("[c d] a, b -> a [e] b", x, y)
```

vectorizes the elementary operation ``{OP}`` according to the expression ``"[c d] a, b -> a [e] b"``.

The meaning of the string expression is defined by analogy with loop notation as follows. The full operation ``einx.{OP}`` will yield the same output as if the elementary
operation ``{OP}`` were invoked in an analogous loop expression:

```python
for a in range(...):
    for b in range(...):
        z[a, :, b] = {OP}(x[:, :, a], y[b])
```

See the [tutorial](https://einx.readthedocs.io/en/latest/gettingstarted/basics.html) for how an einx expression is mapped to the analogous loop expression.

## How are einx operations implemented?

The analogy with loop notation is used only to define what the output of an operation will be. Internally, einx operations are compiled to Python code snippets that invoke operations from the respective tensor framework, rather than using for loops.

The compiled code snippet can be inspected by passing `graph=True` to the einx operation. For example:

```python
>>> x = np.zeros((2, 3))
>>> y = np.zeros((3, 4))
>>> code = einx.add("a b, b c -> c b a", x, y, graph=True)
>>> print(code)

import numpy as np
def op(a, b):
    a = np.transpose(a, (1, 0))
    a = np.reshape(a, (1, 3, 2))
    b = np.transpose(b, (1, 0))
    b = np.reshape(b, (4, 3, 1))
    c = np.add(a, b)
    return c
```

Different [backends](https://einx.readthedocs.io/en/latest/gettingstarted/backends.html) may be used to compile an operation to different implementations, for example following Numpy-like notation, vmap-based notation, or einsum notation.

## Which operations are supported?

**Operations in the API:** einx supports a large set of tensor operations in the namespace ``einx.*``, including reduction, scalar, indexing, some shape-preserving operations, identity map and dot-product. See the [documentation](https://einx.readthedocs.io/en/latest/api/operations.html) for a complete list.

**Operations *not* in the API:** einx additionally allows adapting custom Python functions to einx notation using einx adapters. For example:

```python
# Define a custom elementary operation
def myoperation(x, y):
    x = 2 * x
    z = x + torch.sum(y)
    return z

# Adapt the operation to einx notation
einmyoperation = einx.torch.adapt_with_vmap(myoperation)

# Invoke as einx operation
z = einmyoperation("a [c], b [c] -> a b [c]", x, y)
```

This will yield the same output as if ``myoperation`` were invoked in loop notation:

```python
for a in range(...):
    for b in range(...):
        z[a, b, :] = myoperation(x[a, :], y[b, :])
```

The interface of ``einmyoperation`` matches that of other einx operations. For example, the compiled code snippet can be inspected using ``graph=True``:

```python
>>> code = einmyoperation("a [c], b [c] -> a b [c]", x, y, graph=True)
>>> print(code)
# Constant const1: <function myoperation at 0x49e9aa3cd8a0>
import torch
def op(a, b):
    def c(d, e):
        f = const1(d, e)
        assert isinstance(f, torch.Tensor), "Expected 1st return value of the adapted function to be a tensor"
        assert (tuple(f.shape) == (3,)), "Expected 1st return value of the adapted function to be a tensor with shape (3,)"
        return f
    c = torch.vmap(c, in_dims=(None, 0), out_dims=0)
    c = torch.vmap(c, in_dims=(0, None), out_dims=0)
    g = c(a, b)
    return g
```

See the [documentation](https://einx.readthedocs.io/en/latest/api/adapters.html) for a list of supported adapters.
