einshape

library1.0pypypi✓ verified 90d ago

einshape is a DSL-based reshaping library designed to unify and simplify array manipulation operations such as reshape, squeeze, expand_dims, and transpose, similar to how `einsum` unifies `matmul` and `tensordot`. It primarily targets JAX and TensorFlow frameworks. The current version is 1.0, released in December 2022, indicating a stable but currently infrequent release cadence.

ai-ml

pip install einshape

github ↗package ↗

INSTALL

IMPORT

SIG · EINSHAPE

einshape

ai-mlpythonv1.0

Install

—

Import

—

Disk

—

Pass rate

0/ 10

Env Coverage0 / 10

glibc

3.9–3.13

musl

3.9–3.13

Install & Compatibility

Where this runs

tested against v1.0 · pip install

no network on importno background threads

Install × environment matrix

Each cell = how many times install + import succeeded across repeated harness runs. Partial = flaky.
glibc = Debian/Ubuntu slim · musl = Alpine Linux

musl

glibc

py 3.10

4/8 runs

py 3.11

4/8 runs

py 3.12

4/8 runs

py 3.13

4/8 runs

py 3.9

4/8 runs

Code

Verified usage

Verified import paths — ran on the pinned version, not inferred.

jax_einshape

✓ from einshape import jax_einshape as einshape

This is the primary entry point for JAX users, commonly aliased as 'einshape'.

engine

✓ from einshape import engine

For building custom backend implementations, the parser and engine are exposed.

This quickstart demonstrates basic reshaping operations using `einshape` with JAX. It covers transposing dimensions, combining multiple leading dimensions, and splitting a dimension, highlighting the DSL syntax. Note that JAX must be installed separately for this example to run.

import jax.numpy as jnp from einshape import jax_einshape as einshape x = jnp.arange(2 * 3 * 4).reshape((2, 3, 4)) print(f"Original shape: {x.shape}\n{x}\n") # Equivalent to transpose(x, perm=[0,2,1]) y = einshape("abc->acb", x) print(f"Transposed (abc->acb) shape: {y.shape}\n{y}\n") # Equivalent to reshape combining leading dimensions z = einshape("ab...->(ab)...", x) print(f"Combined leading dims (ab...->(ab)...) shape: {z.shape}\n{z}\n") # Equivalent to splitting a dimension w = einshape("(ab)c->abc", z, a=2) print(f"Split dim ((ab)c->abc) shape: {w.shape}\n{w}")

Debug

Known issues

gotchaeinshape does not list JAX or TensorFlow as direct dependencies. Users must install their preferred array backend (e.g., JAX) separately for `einshape` to be functional for array manipulations.

fix

Ensure `jax` and `jaxlib` (for JAX) or `tensorflow` (for TensorFlow) are installed: `pip install jax jaxlib` or `pip install tensorflow`.

affects: 1.0

gotchaUnderstanding the DSL for grouped dimensions `(components)` and ellipsis `...` is crucial. When splitting a grouped dimension, explicit keyword arguments (e.g., `n=batch_size`) are often required to specify the size of at least one of the new dimensions. Failing to provide these can lead to `ValueError` or incorrect shapes.

fix

Carefully review the `einshape` documentation on DSL syntax, especially for grouped dimensions. For example, `einshape('(mn)hwc->mnhwc', x, n=batch_size)` requires `n` to be specified.

affects: 1.0

gotchaAll index names present on the left-hand side of an `einshape` equation must also be present on the right-hand side, unless they are being implicitly squeezed or combined. Forgetting a dimension or adding an unmatching dimension on the right can lead to shape errors.

fix

Double-check the input and output dimension labels in the `einshape` equation to ensure logical consistency and proper transformation.

affects: 1.0

Errors

Common errors & fixes

ModuleNotFoundError: No module named 'jax'

The JAX library is not installed, but `jax_einshape` was imported or called.

fix

Install JAX and JAXlib: `pip install jax jaxlib`.

ValueError: Axes lengths incompatible: X and Y

The dimensions specified in the `einshape` equation do not allow for a valid reshape operation given the input array's shape, often due to an incorrect grouping, splitting, or missing keyword arguments for new dimensions.

fix

Review the `einshape` equation and the input array's shape. If splitting dimensions, ensure all necessary sizes are provided via `kwargs` (e.g., `einshape('(ab)c->abc', array, a=expected_a_size)`).

ValueError: Equation 'ab->a' is not a valid equation. Every index name that is present on the left-hand side of an equation must also be present on the right-hand side.

Attempted to drop a dimension (e.g., 'b' in 'ab->a') without using a valid `einshape` operation like implicit squeezing (e.g., `a1b->ab`). The DSL requires all LHS indices to be on RHS, unless implied by other transformations.

fix

Ensure all dimensions on the left-hand side are accounted for on the right-hand side. For squeezing, use `1` to denote a unit dimension to be removed, e.g., `a1b->ab` instead of `ab->a`.

Upgrade

Version history

1.0latest on PyPI · released Dec 19, 2022

Audit

Dependencies

absl-pyrequiredCore dependency for logging and utilities.

numpyrequiredCore dependency for array operations.

dataclassesoptionalRequired for Python versions older than 3.7.

jaxoptionalPrimary array backend; required for `jax_einshape`. Install with `pip install jax jaxlib`.

tensorflowoptionalAlternative array backend; required for `tf_einshape` (if exposed).

Agent activity

32 hits · last 30 days

node

Amazon

OpenAI (training)

Resources

githubhttps://github.com/deepmind/einshape ↗

packagehttps://pypi.org/project/einshape/ ↗