See how tensor operations move, reuse and combine elements.
Documentation · Notebook examples · Agent instructions · Releases
rainbow-tensor turns tensor shapes and operations into SVG figures for Jupyter notebooks, lessons and technical explanations. Choose an output element to see its source coordinates and the calculation behind it. Record several operations to follow that element all the way back to the original inputs.
One result, three contributions, two original positions. The first output is
6 + 3 + 6 = 15, because the same source element was sampled twice.
Python 3.10 or newer is required.
python -m pip install rainbow-tensorThe standard installation includes clickable result cells and notebook
controls through ipywidgets and anywidget, along with NumPy and IPython.
Install into the Python environment used by your notebook kernel. PyTorch,
JAX and TensorFlow are optional and installed separately. Live controls need
a notebook host with widget support. Static SVG figures also work in scripts.
The guided examples below target the 1.5.0 source checkout. To try them before
that version is published, run python -m pip install -e . from the repository
root in your notebook's Python environment.
Paste this into a notebook cell:
import numpy as np
import rainbow_tensor as rt
from IPython.display import display
x = np.arange(1, 7).reshape(2, 3)
# [[1, 2, 3],
# [4, 5, 6]]
sum_explorer = rt.explore(rt.sum, x, axis=1, focus=(1,))
display(sum_explorer)
visual = sum_explorer.visualThe result has shape (2,). Output (1,) highlights the second row and
explains 4 + 5 + 6 = 15. Click output (0,) to follow the first row.
The explanation appears below the figure and is also available through
sum_explorer.visual.text. Call sum_explorer.close() when finished.
For a script or a host without live widgets, use rt.sum(x, axis=1, focus=(1,)).
To inspect an indexing expression instead:
selection = ([1, 0, 1], slice(None, None, -1))
display(rt.index(x, selection, focus=(2, 0)))The result is [[6, 5, 4], [3, 2, 1], [6, 5, 4]]. Its position (2, 0) reads
source (1, 2). Repeated picks keep their separate positions in the result.
Flow records the steps that produced a result. Each step retains its shape
and coordinate mapping without allocating a complete intermediate tensor.
flow = rt.Flow()
source = flow.input(x, name="X")
selected = flow.index(source, selection, name="S")
transposed = flow.transpose(selected, name="T")
y = flow.sum(transposed, axis=1, name="Y")
assert y.shape == (3,)
assert [y.value((i,)) for i in range(3)] == [15, 12, 9]
display(y.visualize(focus=(0,)))For Y[0], the trace follows T[0, 0], T[0, 1] and T[0, 2] through the
transpose and index back to X[1, 2], X[0, 2] and X[1, 2]. The duplicate
contribution remains visible.
trace = y.trace((0,))
assert trace.complete
counts = {item.reference.coordinate: item.count for item in trace.roots}
assert counts == {(1, 2): 2, (0, 2): 1}A structural trace reads no array values. Input occurrence counts describe
paths through the recorded expression, not coefficients or derivatives.
Record operations explicitly through flow to retain their history. Input
values are read afresh on each value query or visualization.
Read the provenance guide or run the worked notebook.
Pass a tracked result or an operation and its arguments to explore.
The required widget packages are included in the standard installation:
explorer = rt.explore(y)
display(explorer)
# A single operation can be explored without recording a Flow.
index_explorer = rt.explore(rt.index, x, selection)
display(index_explorer)Click a visible result cell to update the source highlights and explanation. Enter and Space activate a focused cell. Coordinate fields also reach positions hidden by a large preview.
explorer.set_focus((2,))
explorer.visual.save("focused-result.svg")
# Run these when you have finished with the controls.
explorer.close()
index_explorer.close()Live controls need a running notebook kernel and a host with widget support.
y.visualize(focus=(2,)) produces the equivalent static figure.
Use a walkthrough when the learner needs to see a subtotal grow:
lesson = rt.walkthrough(rt.mean, x, axis=1, focus=(1,))
display(lesson)For the row [4, 5, 6], Next term advances through subtotals 4, 9 and 15.
The final mean remains 15 / 3 = 5. A walkthrough also accepts a recorded
Flow result, so the learner can select an intermediate occurrence, such as
the row sum used as a denominator. lesson.snapshot exposes the current
term and numerical state. Call lesson.close() when finished.
To compare reduction axes, start with a shape prediction:
playground = rt.reduction_playground(rt.sum, x, axis=1)
display(playground)Enter (2,), then reveal the answer. Turn on keepdims and predict (2, 1).
The source and matching NumPy expression remain visible while the answer is
hidden. playground.set_parameters(axis=0) applies another recipe from Python.
Call playground.close() when finished.
The guided lessons connect elementwise arithmetic, subtotals and row normalization with executable examples.
| What you want to understand | Public views |
|---|---|
| Tensor structure | shape |
| Slices, masks and repeated gathers | index, take, repeat |
| Reshaping and axis movement | reshape, transpose, swapaxes, moveaxis, squeeze, expand_dims |
| Joining and broadcasting | concatenate, stack, broadcast |
| Elementwise arithmetic | add, subtract, multiply, divide |
| Reductions and contractions | sum, mean, matmul, einsum |
| Storage metadata | memory |
Operation views support focus= to explain one result element. The same
operations are available as Flow methods. shape and memory are standalone
inspection views. flow.broadcast returns one tracked output per input, each
with its own values and output port.
Elementwise views pair the two broadcast source coordinates while retaining
their order. rt.divide(x, 2) accepts a scalar literal. In a Flow, register
that constant with flow.input(2) first. Binary traces expose
trace.expression.operator and ordered trace.expression.operands, keeping
subtraction and division distinct from sums of products.
Indexing supports integers, slices, ellipsis, new axes, boolean masks and
advanced integer arrays. Reductions support multiple axes, negative axes,
axis=None, axis=() and keepdims=True. Scalars use shape and coordinate
(). Empty tensors retain their zero-length dimensions and have no selectable
output element.
NumPy arrays and CPU tensors from PyTorch, JAX and TensorFlow are covered by
the backend contract suite. Other array-like inputs need a shape and scalar
coordinate access. Shape tuples such as rt.shape((2, 3, 4)) use generated
row-major placeholder values.
Every view returns a TensorVisual with an SVG, a plain-text explanation and
operation metadata. Its trace describes a focused output. Visuals from a
recorded Flow also expose a bounded provenance tree.
visual.save("row-sum.svg")
print(visual.text)save writes the figure only. Keep visual.text alongside it when sharing the
explanation outside a notebook. SVG figures remain usable without a running
Python process.
Large tensors use bounded previews. Repeat mappings and basic slice selections
remain compact instead of enumerating every output position. Math views default
to 10,000 terms per output and 100,000 across the preview. Flow evaluation also
counts recursive factor references and input reads. Work that exceeds a budget
is shown as ?, and incomplete structural traces identify the limit reached.
Numerical previews use Python scalar arithmetic. Backend accumulation dtype,
rounding and overflow may differ. The figures explain logical operations and
coordinate origins, while memory reports available storage metadata.
Division by a zero denominator raises ZeroDivisionError, including 0 / 0.
The default auto theme follows the viewer's light or dark preference, including
in saved SVG files. Language follows the Python environment and system locale.
English and Simplified Chinese are included.
rt.set_default_theme("auto")
rt.set_language("en")
display(rt.shape(x, theme="dark")) # Override the theme for one figure.Add a language by supplying a JSON translation catalog. Missing entries fall
back to English, and rt.load_translations("translations") loads extra catalogs
without changing application code. See
themes
and translations.
| Start here | What you will find |
|---|---|
| Learning path | Small exercises that connect shapes, output coordinates and source values |
| Verified lessons | Canonical teaching examples whose assertions and focus changes run in CI |
| Guided lessons | Binary arithmetic, contribution walkthroughs, row normalization and axis prediction |
| Notebook collection | Runnable examples from basic shapes through cross-operation tracing |
| API reference | Function signatures, parameters and result objects |
| Agent instructions | Instructions and runnable patterns for agents generating visual NumPy explanations |
| Architecture | Coordinate mappings, rendering boundaries and performance decisions |
Give the agent instructions URL to your agent with a NumPy question. The page teaches the agent how to generate runnable Rainbow Tensor code and explain the result visually.
Bug reports, clearer teaching examples and translation catalogs are welcome. For a bug report, include a small reproducing input, the operation, the expected result and your package, Python and backend versions. Use GitHub Issues to report a problem or discuss a new operation.
To work on the package:
git clone https://github.com/Niox1337/rainbow-tensor.git
cd rainbow-tensor
python -m pip install -e ".[dev]"
python -m pytest
python -m ruff check .
python scripts/check_lessons.pyOptional backend tests skip when their framework is absent from that Python environment. CI separately requires PyTorch, JAX and TensorFlow, as well as running the main suite on Python 3.10 and 3.12. SVG golden tests protect existing figures, and distribution checks exercise installed wheels and source packages.
python -m build
python scripts/check_distribution.py
python -m pip install -r docs/requirements.txt
python -m sphinx -b html -W --keep-going docs docs/_build/htmlThe distribution check expects one wheel and one source archive in dist.
Use --dist-dir PATH for a different build directory. Keep changes focused and
add a regression example when correcting an operation's behavior.
MIT, Copyright 2026 Zhixiang Feng.