Colourful SVG visuals for tensor shapes, indexing, and operations, built for Jupyter notebooks and teaching.
rt.shape(np.arange(8).reshape(2, 2, 2))
Tensor shapes are hard to hold in your head, and an index like (0, slice(None), 1) gives no hint of what it selects until you run it. A printed array is just a wall of numbers. rainbow-tensor draws the tensor as nested coloured frames and highlights exactly which elements an operation touches, so the structure and the result are clear at a glance.
Every axis keeps one colour through every view, so you can follow an axis as it moves, folds, or stretches. That makes it a fast way to learn how reshapes and reductions work, to teach shape transformations, and to debug a confusing indexing or broadcasting bug. The core imports no deep learning framework, so it stays light and works with plain NumPy or any array that exposes a shape.
pip install rainbow-tensorThe distribution name is rainbow-tensor and the import name is rainbow_tensor.
Run inside a Jupyter notebook or an IPython shell so the SVG is displayed. The convention is to import the package as rt.
import numpy as np
import rainbow_tensor as rt
x = np.arange(8).reshape(2, 2, 2)
rt.shape(x) # draw the structure
rt.index(x, (0, slice(None), 1)) # highlight what an index selects
Each call returns a small result object. Its svg attribute holds the SVG string, its text attribute holds the explanation printed under the figure, and save writes the SVG to a file.
Shape changing, combining, and broadcasting views draw the source and the result side by side, so the mapping between them is easy to follow.
- Shapes and indexing with
shapeandindex, covering integers, slices, ellipsis, new axes, boolean masks, and fancy integer arrays - Reshaping and moving axes with
reshape,transpose,swapaxes,moveaxis,squeeze, andexpand_dims - Reductions and math with
sum,mean,matmul, andeinsum - Combining with
concatenate,stack,broadcast,repeat, andtake
rt.sum(np.arange(12).reshape(3, 4), 0)
rt.einsum("ij,jk->ik", np.arange(6).reshape(2, 3), np.arange(12).reshape(3, 4))
Pass theme="dark" to any call, or set a module default that every later call follows. A theme is a plain object you can tweak with variant, and a global axis ramp can be set once with set_default_axis_colors.
rt.shape(x, theme="dark")
rt.set_default_theme("dark")The full guide and API reference live at rainbow-tensor.zhixiangfeng.com.
Runnable notebooks for every feature group live in examples, and more sample images live in examples/images.
pip install -e ".[dev]"
pytest
ruff check .
python -m buildMIT