Einshape
A DSL-based library for unified tensor reshaping, squeezing, expanding, and transposing in JAX, TensorFlow, and NumPy.
What is Einshape?
Einshape is a Python library that provides a domain-specific language (DSL) for tensor reshaping operations, similar to how `einsum` standardizes tensor contractions. It unifies operations like reshape, squeeze, expand_dims, and transpose under a single equation-based syntax, making complex array manipulations more intuitive and less error-prone.
Machine learning researchers, data scientists, and engineers working with JAX, TensorFlow, or NumPy who frequently manipulate tensor shapes in their workflows.
Developers choose Einshape for its expressive, concise DSL that reduces boilerplate code and mental overhead compared to chaining multiple native reshaping functions, while maintaining compatibility across major numerical frameworks.
Overview
Einshape is a domain-specific language (DSL) library that provides a unified interface for tensor reshaping operations across multiple frameworks. It extends the intuitive equation syntax of einsum to operations like reshape, squeeze, expand_dims, and transpose, making tensor manipulations more expressive and consistent.
Key Features
- Unified Reshaping DSL — Uses
lhs->rhsequations to specify complex tensor transformations concisely. - Multi-Framework Support — Works with JAX, TensorFlow, and NumPy via separate implementations.
- Advanced Axis Manipulation — Supports ellipsis (
...) for arbitrary ranks, groups( )for splitting/combining dimensions, and units (1) for squeezing/expanding. - Dynamic Shape Control — Allows runtime shape specifications via keyword arguments (e.g.,
n=batch_size). - Consistent Syntax — Mirrors
einsum's approach, reducing cognitive load for users familiar with Einstein summation.
Philosophy
Einshape aims to bring the clarity and power of DSL-based tensor operations beyond contraction to reshaping, promoting a more declarative and readable style of array manipulation.
Related Projects
Found a gem we're missing?
Open-Awesome is built by the community, for the community. Submit a project, suggest an awesome list, or help improve the catalog on GitHub.
