A collection of Manim Community Edition animations for visualizing einops operations.

• einops_manim_lib.py - Common library with shared visualization functions
• einsum_generalised.py - Animation for einops.einsum() operation
• reduce_anim.py - Animation for einops.reduce() operation
• repeat_anim.py - Animation for einops.repeat() operation

python3 -m venv .venv
source .venv/bin/activate
uv sync

Edit the configuration in einsum_generalised.py:

TENSOR_NAMES = ["a", "b"]
EINSPEC = "h c w, h c w -> w"

Run:

docker run --rm -it -v "./einops-anim:/manim" manimcommunity/manim manim -qm einsum_generalised.py EinsumBudgeted

Edit the configuration in reduce_anim.py:

TENSOR_NAME = "x"
PATTERN = "h w c -> h w"  # Reduce over channels
REDUCTION = "mean"  # Options: 'mean', 'sum', 'max', 'min', 'prod'

Examples:

• "h w c -> h w" with REDUCTION = "mean" - Average pooling over channels
• "b h w c -> b c" with REDUCTION = "max" - Max pooling over spatial dimensions
• "h w -> h" with REDUCTION = "sum" - Sum over width

Run:

docker run --rm -it -v "./einops-anim:/manim" manimcommunity/manim manim -qm reduce_anim.py ReduceAnim

Edit the configuration in repeat_anim.py:

TENSOR_NAME = "x"
PATTERN = "h w -> h w c"  # Add new axis
REPEAT_PARAMS = {"c": 3}  # Repeat 3 times along new axis

Examples:

• "h w -> h w c" with REPEAT_PARAMS = {"c": 3} - Add new axis with repetition
• "h w -> (repeat h) w" with REPEAT_PARAMS = {"repeat": 2} - Repeat along existing axis
• "h w -> h w 1" - Add singleton dimension

Run:

docker run --rm -it -v "./einops-anim:/manim" manimcommunity/manim manim -qm repeat_anim.py RepeatAnim

• Up to 4D tensors - Visualizes 2D, 3D, and 4D tensors with appropriate layouts
• Automatic sizing - Axes are automatically assigned sizes (2 or 3) for compact display
• Slow-paced animations - Increased run times for better comprehension
• Centered computation display - Arithmetic shown in center of screen between tensors
• Element tracking - Shows how individual elements transform between input and output

• Shows all reduction terms (no budgeting)
• Displays tensor[indices]=value format for clarity
• Highlights corresponding cells in input tensors during multiplication
• Step-by-step accumulation of sum

• Highlights all input elements being reduced
• Shows reduction formula (mean, sum, max, etc.)
• Updates output cells with computed values

• Shows which input elements are copied to multiple output locations
• Highlights repeated values across output tensor

All animations support these common settings:

RNG_SEED = 42  # Random seed for reproducible tensors
AXIS_SIZE_CHOICES = (2, 3)  # Possible sizes for auto-assigned axes

The einops_manim_lib.py provides reusable functions:

• grid_from_2d() - Display 2D array as grid
• grid_from_3d() - Display 3D tensor as stacked grids
• grid_from_4d() - Display 4D tensor as nested grids
• display_tensor() - Automatic display based on dimensionality

• highlight_cell() - Highlight 2D grid cell
• highlight_cell_3d() - Highlight 3D grid cell
• highlight_cell_4d() - Highlight 4D grid cell

• safe_text() - Create Text objects (no LaTeX)
• parse_axes_spec() - Parse axis specifications
• assign_minimal_axis_sizes() - Auto-assign axis sizes

To create animations for other operations:

1. Import the library: from einops_manim_lib import *
2. Use display_tensor() to show tensors
3. Use highlighting functions to show element correspondences
4. Follow the layout pattern: header at top, input left, output right, computation center

Error: "Missing $ inserted"

• The code avoids LaTeX by using Text() instead of Tex() or MathTex()

Tensors too small to see

• Adjust AXIS_SIZE_CHOICES to use larger values like (3, 4)
• Modify cell_size and font_size parameters in display functions

Animation too fast

• Increase run_time parameters in self.play() calls
• Increase self.wait() durations

Type errors with indices

• Check that tensor dimensions match the pattern specification
• Verify AXIS_PARAMS or REPEAT_PARAMS are correctly set

This code is provided as-is for educational purposes.