einx

What it is and what it does

einx is a notation system and compiler that lets you write tensor operations once using a string-based syntax and run them across NumPy, PyTorch, JAX, TensorFlow, and MLX without rewriting code. The notation works by analogy to nested loops: you describe how an elementary operation (like addition or dot product) should be vectorized across tensor axes using bracket and parenthesis syntax, and einx compiles that into optimized backend-specific code.

The library supports a large set of built-in operations (reductions, scalar ops, indexing, reshaping, dot products) and also lets you wrap custom Python functions as einx operations using framework-specific adapters like vmap. It's in active development (Alpha status) with low install friction and no known security vulnerabilities.

Use it for:

• Write portable tensor code that runs on NumPy for CPU prototyping and PyTorch or JAX for GPU training without duplicating logic.
• Express complex reshaping, pooling, and gather operations using a single notation instead of chaining framework-specific calls.
• Adapt custom functions (e.g., domain-specific kernels) to einx notation to gain automatic vectorization across batch and spatial dimensions.
• Compile tensor operations to inspectable code snippets for debugging or optimization verification.
• Build machine learning pipelines that can switch backends at runtime by changing only the framework import.

Worth the install?

AI-flagged interpretation of the facts on this page — verify before relying

einx provides a universal notation for expressing tensor operations across multiple frameworks (NumPy, PyTorch, JAX, TensorFlow, MLX) using a single string-based syntax that compiles to optimized code for each backend.

Yes, if you work across multiple tensor frameworks or want a more expressive notation for complex tensor operations. The low install friction, active maintenance, MIT license, and zero known vulnerabilities make it safe to adopt. Start with the quickstart if you're unsure whether the notation fits your workflow—it's most valuable for teams managing multi-backend code or expressing operations that are verbose in native framework syntax.

Install

einx on PyPI

pip install einx

uv add einx

poetry add einx

Installing einx

Before you install

Installation is straightforward with low friction—a pure Python wheel with three stable runtime dependencies (numpy, sympy, frozendict). The package is actively maintained with a recent release and steady commit activity.

License in practice

MIT license permits unrestricted use, modification, and distribution in both open-source and commercial projects with minimal obligations.

Quickstart

pip install einx

import einx
import numpy as np

x = np.ones((10, 20, 30))
y = einx.sum("a [b] c", x)  # Sum-reduction along second axis
print(y.shape)

Requires Python 3.10 or later.

Verify before relying

• Performance characteristics and compilation overhead compared to direct framework calls for simple operations.
• Completeness of operation coverage across all five supported backends (NumPy, PyTorch, JAX, TensorFlow, MLX).
• Whether custom adapters work equally well across all backends or have backend-specific limitations.

Package facts

Evidence: einx-0.4.3-py3-none-any.whl

Tags

More Scientific/Engineering packages