arogozhnikov/einops
A tensor operations library that provides flexible and readable tensor manipulation APIs across PyTorch, JAX, TensorFlow, NumPy, and other array computing frameworks.

Einops provides composable tensor operations designed to replace verbose reshape, permute, and reduce operations common in deep learning code. It introduces a readable notation for tensor transformations that works consistently across multiple backends including PyTorch, JAX, TensorFlow, NumPy, MLX, and others. The library includes utilities for packing and unpacking tensors as well as einsum operations.