einx - GitLite

KO | EN

gitlite — search

#javascript #react #typescript #llm #go #golang #rust #linux #docker #java #windows #awesome #security #mcp

einx

(★ 525)

Universal Notation for Tensor Operations in Python

파일 탐색기

최종 버전 다운로드 (.zip)

- - publish_pypi.yml
  - run_pytest.yml
  - run_typecheckers.yml
- - conf.py
  - generate_compiled_code.py
  - index.rst
- make.bat
- Makefile
- requirements.txt
- - __init__.py
  - functorchdim.py
- __init__.py
- arrayapi.py
- backend.py
- errors.py
- jax.py
- mlx.py
- nn.py
- numpy.py
- py.typed
- tensorflow.py
- tinygrad.py
- torch.py
- conftest.py
- test_adapt.py
- test_backend_resolution.py
- test_compare_libs.py
- test_compilation.py
- test_experimental.py
- test_invalid_backend.py
- test_invalid_backend_version.py
- test_shapes.py
- test_tensor_conversion.py
- test_values.py
- fail.py
- pass.py
.gitignore
.pre-commit-config.yaml
.readthedocs.yml
CHANGELOG.md
LICENSE
pyproject.toml
README.md
ruff.toml

# 설치 가이드

1. 코드 내려받기

git clone https://github.com/fferflo/einx

깃허브에서 프로젝트 코드 전체를 내 컴퓨터로 내려받습니다.

cd einx

방금 내려받은 프로젝트 폴더 안으로 이동합니다.

2. 공식 설치 스크립트

쉬움 추천

사전 준비물

Python 3 pip 명령어를 쓰려면 Python이 필요합니다.

pip install einx

PyPI에 배포된 패키지를 바로 설치합니다. 소스 클론이 필요 없습니다.

✅ 설치 후 새 터미널을 열고, 프로그램의 버전 확인 명령(예: --version)으로 정상 설치됐는지 확인하세요.

이 레포의 README에 적힌 실제 명령어를 그대로 가져왔습니다.

3. Python

쉬움

사전 준비물

Git GitHub에서 프로젝트 코드를 내려받으려면 필요합니다.
Python 3 설치 시 'Add Python to PATH' 옵션을 꼭 체크하세요 (Windows).

pip install einx

PyPI에 배포된 패키지를 바로 설치합니다. 소스 클론이 필요 없습니다.

✅ 에러 메시지 없이 실행되고 터미널에 안내 문구가 출력되면 정상입니다.

이 레포의 README에 적힌 실제 명령어를 그대로 가져왔습니다.

// repository documentation

Was this content helpful?

★ 0 (0 ratings)

# *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://img.shields.io/pypi/v/einx.svg?color=blue)](https://pypi.org/project/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)

*Update:* einx has been accepted at ICLR 2026 as an oral ([paper](https://openreview.net/pdf?id=QqvQ3iAdpC)).

## 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.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
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
```

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?

**Built-in operations:** 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.

**User-defined operations:** einx additionally allows adapting custom Python functions to einx notation using einx adapters. For example:

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

# Adapt the operation to einx notation
einfoo = einx.torch.adapt_with_vmap(foo)

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

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

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

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

```python
>>> code = einfoo("a [c], b [c] -> a b [c]", x, y, graph=True)
>>> print(code)
# Constant const1: <function foo 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.