rainbow-tensor

Visualise tensor shape, indexing, and slicing as SVG inside IPython and Jupyter notebooks.

rainbow-tensor is made for people who are learning how a tensor is structured and how an indexing expression selects elements. It draws the tensor as nested blocks, rows, and cells, then highlights exactly which elements an index picks out.

Examples

rt.shape(np.arange(8).reshape(2, 2, 2))

rt.index(np.arange(8).reshape(2, 2, 2), (0, slice(None), 1))

The same view also renders in a dark theme.

rt.shape(np.arange(8).reshape(2, 2, 2), theme="dark")

Shape changing, combining, and broadcasting views draw the source and the result side by side.

rt.reshape((2, 3), (3, 2))

rt.transpose((2, 3))

rt.sum((3, 4), 0)

rt.mean((3, 4), 1)

rt.concatenate([(2, 3), (2, 3)], 0)

rt.stack([(2, 3), (2, 3)], 0)

rt.broadcast((3, 1), (1, 4))

rt.einsum("ij,jk->ik", np.arange(6).reshape(2, 3), np.arange(12).reshape(3, 4))

rt.swapaxes(np.arange(24).reshape(2, 3, 4), 0, 2)

The same calls work with backend arrays that expose a shape and coordinate access.

rt.shape(np.arange(6).reshape(2, 3))

More sample images live in examples/images, and runnable notebooks live in examples.

Features

• Static SVG output that stays sharp at any zoom level in a notebook
• Shape visualisation for tensors of any rank, nesting frames to arbitrary depth
• Index visualisation with highlighted selections and a plain text explanation, including boolean masks and integer array indexing
• Shape changing views for reshape, transpose, and axis reductions, drawing the source and the result side by side
• Swapaxes views that move two axes while keeping the source colours traceable
• Combining views for concatenate and stack, tinting each operand so the seam or the new axis is clear
• Broadcasting views that stretch a smaller operand to match a larger one, marking every stretched axis
• Einsum views that colour shared labels, mark contracted axes, and show the derived output shape
• Backend value access for NumPy style arrays plus Torch, JAX, and TensorFlow style scalar values
• A renderer registry with SVG as the default output backend
• A light theme and a dark theme, selectable per call or through a module default
• An axis legend that names each axis with its size in the matching colour
• Configurable float precision with right aligned numbers
• Long axes truncate to a readable head and tail with an ellipsis cell
• Hover any cell to read its coordinate and flat index
• A save helper that writes the SVG to a file
• Works with shape tuples and with array-like objects that expose a .shape attribute, such as NumPy arrays
• No tensor library is imported by the core, so the package stays lightweight

Colour scheme

Each axis has its own colour drawn from a rainbow ramp keyed by depth, so the structure and a selection are easy to read.

• Axis 0 is the outer frame, drawn red
• Axis 1 is the inner row frame, drawn orange
• Deeper axes continue through amber, green, blue, violet, and pink
• The leaf axis elements sit in plain cells, and a selected element fills green
• The numbers in the shape label, the tokens in the index label, and the legend swatches are coloured to match

In an index view only the selected frames keep their axis colour. The rest of the tensor is dimmed so the selected path stands out.

Themes

Pass theme="light" or theme="dark" to any call, or set a module default that every later call follows.

import numpy as np
import rainbow_tensor as rt

x = np.arange(8).reshape(2, 2, 2)
rt.shape(x, theme="dark")

rt.set_default_theme("dark")
rt.index(x, (0, slice(None), 1))

A theme bundles the colours, fonts, cell size, stroke width, and the truncation limit. Derive a tweaked copy with variant and pass it directly.

roomy = rt.LIGHT.variant(cell_w=64, max_cells=8)
rt.shape(np.arange(8).reshape(2, 4), theme=roomy)

Float precision and saving

Control how floats are formatted with precision, then write the SVG to a file with save.

import numpy as np
import rainbow_tensor as rt

x = np.linspace(0, 1, 6).reshape(2, 3)
visual = rt.shape(x, precision=3)
visual.save("tensor.svg")

Shape changing operations

Beyond viewing a single tensor, rainbow-tensor draws the operations that rearrange or summarise one. Each view places the source and the result in one figure with a connector, so the mapping is easy to follow.

import numpy as np
import rainbow_tensor as rt

x = np.arange(12).reshape(3, 4)

rt.reshape(x, (2, 6))      # the same values flow into a new layout
rt.transpose(x)            # axes reverse, each keeping its colour
rt.transpose(x, (1, 0))    # an explicit permutation
rt.swapaxes(x.reshape(2, 3, 2), 0, 2)   # swap exactly two axes
rt.sum(x, 0)               # collapse axis 0, the result keeps the rest
rt.mean(x, 1)              # collapse axis 1 into per group means

reshape keeps the row major order, so element k stays element k. A single -1 lets one axis be inferred. transpose colours each result axis by the source axis it came from, so a colour can be traced across the move. swapaxes does the same for two chosen axes. sum and mean mark the source elements that fold into the first result element and draw the surviving shape.

Combining tensors

concatenate and stack join several operands into one. Each operand is drawn in its own tint, and the result colours every cell by the operand it came from, so the seam between operands or the new axis stays clear.

import numpy as np
import rainbow_tensor as rt

a = np.arange(6).reshape(2, 3)
b = np.arange(100, 106).reshape(2, 3)

rt.concatenate([a, b], 0)   # join along an existing axis, (4, 3)
rt.concatenate([a, b], 1)   # grow the columns instead, (2, 6)
rt.stack([a, b], 0)         # place onto a new leading axis, (2, 2, 3)

concatenate needs the operands to match on every axis except the joined one, while stack needs them to share one shape. A mismatch raises a clear error rather than drawing a wrong figure.

Broadcasting

broadcast stretches a smaller operand to match a larger one. Each operand is drawn in its own shape and again stretched to the common broadcast shape, so the repeated values along a stretched axis are visible, and every stretched axis is marked in the accent colour.

import numpy as np
import rainbow_tensor as rt

a = np.arange(3).reshape(3, 1)
b = np.arange(4).reshape(1, 4)

rt.broadcast(a, b)          # (3, 1) and (1, 4) stretch to (3, 4)
rt.broadcast((2, 3, 4), (4,))   # a (4,) vector gains two leading axes

Axes line up from the right, and on each axis the sizes must be equal or one of them must be 1. Incompatible shapes raise a clear error rather than drawing a wrong figure.

Einsum

einsum turns a subscript expression into labelled operand panels and an output panel. A repeated label that does not appear in the output is highlighted on the operands, which makes the contraction visible.

import numpy as np
import rainbow_tensor as rt

a = np.arange(6).reshape(2, 3)
b = np.arange(12).reshape(3, 4)

rt.einsum("ij,jk->ik", a, b)
rt.einsum("abc,cde,ef->abdf", (2, 2, 2), (2, 2, 2), (2, 2))

Shared labels keep the same colour wherever they appear. The output shape is derived from the free labels in the output subscript and checked with the same size rules as NumPy.

Backend arrays

rainbow-tensor keeps using duck typing for arrays. If an object exposes a .shape attribute and supports coordinate reads, the visualiser can draw its values without importing that backend in the core package.

import rainbow_tensor as rt
import torch
import jax.numpy as jnp

rt.shape(torch.arange(6).reshape(2, 3))
rt.shape(jnp.arange(6).reshape(2, 3))

Torch, JAX, and TensorFlow checks are optional in the test suite. They run when those packages are installed and skip cleanly otherwise.

Renderer registry

SVG is still the default renderer. A custom renderer can be registered for experiments with other output formats. It receives the same tensor shape, panel dictionaries, value functions, theme, and precision that the SVG renderer uses.

import rainbow_tensor as rt


class TextRenderer:
    name = "text"
    mime_type = "text/plain"

    def render_tensor(self, **kwargs):
        return str(kwargs["shape"])

    def render_panels(self, **kwargs):
        return str(kwargs["panels"][-1]["shape"])


rt.register_renderer(TextRenderer())
rt.shape((2, 3), renderer="text")

Use set_default_renderer to choose a renderer for later calls, or pass renderer= on one call when you only want a local override.

Installation

Install from PyPI.

pip install rainbow-tensor

Install from source for development.

git clone https://github.com/Niox1337/rainbow-tensor.git
cd rainbow-tensor
pip install -e .

Install with the development tools (pytest, ruff, build).

pip install -e ".[dev]"

The distribution name is rainbow-tensor and the import name is rainbow_tensor.

Usage

Run the examples in a Jupyter notebook or an IPython shell so the SVG is displayed.

The convention is to import the package as rt.

Visualise a shape.

import numpy as np
import rainbow_tensor as rt

x = np.arange(8).reshape(2, 2, 2)
rt.shape(x)

Visualise how an index selects elements.

rt.index(x, (0, slice(None), 1))

For the array np.arange(8).reshape(2, 2, 2) the index (0, slice(None), 1) selects the values 1 and 3, the selected coordinates are (0, 0, 1) and (0, 1, 1), and the result shape is (2,).

Each function returns a small result object. Its svg attribute holds the SVG string, so the package can be inspected and tested outside a notebook.

Supported

• Tensors of any rank, with frames nested to arbitrary depth
• Shape tuples and array-like objects with a .shape attribute
• Integer indexing, including negatives such as -1
• Basic slicing with slice(None), slice(start, stop), and slice(start, stop, step), including negative bounds and steps such as slice(None, None, -1)
• Ellipsis (...) to fill the remaining axes, such as (0, ..., 1)
• None (newaxis) to insert a size 1 axis, shown in the result shape and label
• A full-shape boolean mask, highlighting every True position
• Integer array (fancy) indexing, including mixed with slices, with the gathered axis placed as NumPy does
• Reshape with row major order and one inferred -1 axis
• Transpose and permute with axis colours following the move
• Swapaxes with negative axes resolved like NumPy
• Sum and mean reductions over a chosen axis
• Concatenate along an existing axis, with the seam tinted
• Stack onto a brand new axis
• Broadcasting two tensors to a common shape, marking every stretched axis
• Einsum with explicit or implicit output labels and any number of operands
• Backend arrays with .shape and coordinate access, including optional Torch, JAX, and TensorFlow checks
• Custom renderers through the renderer registry

Not supported yet

• Multi-dimensional index arrays and per-axis boolean arrays
• Einsum ellipsis notation
• Interactive controls and animation

Development

pytest
ruff check .
python -m build

License

MIT