{"id":7201,"library":"einshape","title":"einshape","description":"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.","status":"active","version":"1.0","language":"en","source_language":"en","source_url":"https://github.com/deepmind/einshape","tags":["JAX","TensorFlow","array manipulation","deep learning","reshaping","einsum-like"],"install":[{"cmd":"pip install einshape","lang":"bash","label":"Stable release from PyPI"},{"cmd":"pip install git+https://github.com/deepmind/einshape","lang":"bash","label":"Latest from GitHub"}],"dependencies":[{"reason":"Core dependency for logging and utilities.","package":"absl-py","optional":false},{"reason":"Core dependency for array operations.","package":"numpy","optional":false},{"reason":"Required for Python versions older than 3.7.","package":"dataclasses","optional":true},{"reason":"Primary array backend; required for `jax_einshape`. Install with `pip install jax jaxlib`.","package":"jax","optional":true},{"reason":"Alternative array backend; required for `tf_einshape` (if exposed).","package":"tensorflow","optional":true}],"imports":[{"note":"This is the primary entry point for JAX users, commonly aliased as 'einshape'.","symbol":"jax_einshape","correct":"from einshape import jax_einshape as einshape"},{"note":"For building custom backend implementations, the parser and engine are exposed.","symbol":"engine","correct":"from einshape import engine"}],"quickstart":{"code":"import jax.numpy as jnp\nfrom einshape import jax_einshape as einshape\n\nx = jnp.arange(2 * 3 * 4).reshape((2, 3, 4))\nprint(f\"Original shape: {x.shape}\\n{x}\\n\")\n\n# Equivalent to transpose(x, perm=[0,2,1])\ny = einshape(\"abc->acb\", x)\nprint(f\"Transposed (abc->acb) shape: {y.shape}\\n{y}\\n\")\n\n# Equivalent to reshape combining leading dimensions\nz = einshape(\"ab...->(ab)...\", x)\nprint(f\"Combined leading dims (ab...->(ab)...) shape: {z.shape}\\n{z}\\n\")\n\n# Equivalent to splitting a dimension\nw = einshape(\"(ab)c->abc\", z, a=2)\nprint(f\"Split dim ((ab)c->abc) shape: {w.shape}\\n{w}\")","lang":"python","description":"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."},"warnings":[{"fix":"Ensure `jax` and `jaxlib` (for JAX) or `tensorflow` (for TensorFlow) are installed: `pip install jax jaxlib` or `pip install tensorflow`.","message":"einshape 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.","severity":"gotcha","affected_versions":"1.0"},{"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.","message":"Understanding 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.","severity":"gotcha","affected_versions":"1.0"},{"fix":"Double-check the input and output dimension labels in the `einshape` equation to ensure logical consistency and proper transformation.","message":"All 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.","severity":"gotcha","affected_versions":"1.0"}],"env_vars":null,"last_verified":"2026-04-16T00:00:00.000Z","next_check":"2026-07-15T00:00:00.000Z","problems":[{"fix":"Install JAX and JAXlib: `pip install jax jaxlib`.","cause":"The JAX library is not installed, but `jax_einshape` was imported or called.","error":"ModuleNotFoundError: No module named 'jax'"},{"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)`).","cause":"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.","error":"ValueError: Axes lengths incompatible: X and Y"},{"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`.","cause":"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.","error":"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."}]}