rainbow-tensor

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.

Installation

Python 3.10 or newer is required.

python -m pip install rainbow-tensor

The 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 target rainbow-tensor 1.8.0. Install the package in the same Python environment as your notebook kernel.

Your first visualization

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.visual

The 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.

Follow a complete operation chain

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.

Explore by clicking

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.

Build a calculation one contribution at a time

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.

Compare, select and find extrema

Flow keeps a comparison and its conditional choice as separate operations. The structural trace shows possible candidates. The evaluated choice is shown with a thick dashed outline, so learners can distinguish the winning source from the other dependencies.

flow = rt.Flow()
scores = flow.input(np.array([2, 8, 8]), name="Scores")
threshold = flow.input(4, name="Threshold")
condition = flow.greater(scores, threshold, name="Above threshold")
filtered = flow.where(condition, scores, flow.input(0, name="Zero"), name="Filtered")
display(filtered.visualize(focus=(1,)))

position = flow.argmax(scores, name="Best position")
lesson = rt.walkthrough(position)
display(lesson)
assert lesson.snapshot.selected_source.source.coordinate == (1,)
lesson.close()

where preserves condition, true branch, and false branch as distinct roles. For a tie in [2, 8, 8], argmax returns position 1 and marks the first maximum as its source. See the conditions and extrema guide.

Share a portable lesson

Capture states that you have already prepared, then save them for offline playback. Recipients need only a browser.

lesson = rt.explore(rt.sum, x, axis=1)
recording = rt.capture_lesson([lesson])
recording.save("row-sums.html")
lesson.close()

The recording contains only the states you captured. It does not run Python or calculate uncaptured results. See the portable lessons guide.

Supported operations

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.

Figures, explanations and limits

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.

Theme and language

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.

Learn and teach

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.

Contributing

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.py

Optional 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/html

The 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.

License

MIT, Copyright 2026 Zhixiang Feng.