einx is a Python library providing a unified interface for tensor operations across NumPy, PyTorch, JAX, and TensorFlow, using an Einstein-inspired notation. It aims to simplify complex tensor manipulations, particularly for neural network development, by offering a consistent and composable syntax for operations like summation, reshaping, and matrix multiplication.

How It Works

einx employs a novel notation inspired by einops but introduces bracket notation and full composability for a universal tensor language. Operations are defined using this notation, which is then just-in-time compiled into optimized backend-specific Python functions using exec(). This approach allows for seamless integration with existing codebases and efficient execution by reducing overhead through caching.

Quick Start & Requirements

• Install via pip: pip install einx
• Requires Python. Supports NumPy, PyTorch, JAX, and TensorFlow backends.
• See Installation for more details.

Highlighted Details

• Universal notation for tensor operations across multiple frameworks.
• Supports advanced operations like grouped and spatial matrix multiplications.
• Integrates deep learning modules for PyTorch, Flax, Haiku, Equinox, and Keras.
• Just-in-time compilation for optimized performance and inspectable code generation.

Maintenance & Community

• Actively maintained by the einx team.
• Examples and tutorials are available for various deep learning models and frameworks.
• See examples/train_torch.py for example training.

Licensing & Compatibility

• MIT License. Permissive for commercial use and integration with closed-source projects.

Limitations & Caveats

• The library relies on Python's exec() for JIT compilation, which might have security implications in untrusted environments.
• While aiming for universality, the complexity of the notation might require a learning curve for new users.