tensordiagram 0.1.2

Visualize torch, jax, tensorflow, etc. tensors for understanding and debugging.

tensordiagram

A python library for visualizing tensors from torch, jax, tensorflow, numpy, etc. Helps with learning and debugging in notebooks and other contexts. It's built on top of the graphics backend, chalk.

why

Debugging deep learning code is hard—especially when it's foreign, because it's hard to imagine tensor manipulations, e.g. F.conv2d(x.unsqueeze(1), w.transpose(-1, -2)).squeeze().view(B, L, -1) in your head. Printing shapes and tensor values only gets you so far. tensordiagram lets me easily represent tensors visually, inside python code, notebooks, and interpreter sessions.

Other python libraries for creating tensor diagrams are either too physics and math focused, not notebook-friendly, limited to visualizing single tensors, and/or serve a wider purpose (so have a steep learning curve).

install

pip install tensordiagram

Separately, you'll need to install cairo for png images:

# might need to install the base library first on debian
sudo apt-get install libcairo2-dev

# cairo python bindings
pip install ".[cairo]"

Or, for svg images:

pip install ".[svg]"

intro

For an in-depth guide, refer to this notebook.

import numpy as np
import tensordiagram as td

tensor = np.arange(12).reshape(3, 4)
diagram = td.to_diagram(tensor)

The diagram can be saved using render_png or render_svg:

diagram.render_png("output.png", height=300)

Style and annotate diagrams:

diagram \
  .fill_values() \
  .fill_color("lightblue") \
  .annotate_dim_size(dim="row", color="lightgreen")

3d tensor:

papaya = "#ff9700"

tensor = np.arange(24).reshape((2, 3, 4))
diagram = td.to_diagram(tensor).fill_color(papaya)

Combine tensor and chalk diagrams for intricate outputs:

import random
import torch

color_names = [ "red", "blue", "green", "purple", "orange", "pink", "cyan"]

def random_colors_tensor(shape):
  colors_array = np.empty(shape, dtype=object)
  for index, _ in np.ndenumerate(colors_array):
    colors_array[index] = random.choice(color_names)
  return colors_array

color_tensor = random_colors_tensor(shape=(2, 3, 4))

# tensors
t = torch.arange(24).reshape((2, 3, 4))
slice_1 = t[0, :, :].unsqueeze(0)
slice_2 = t[:, 1, :].unsqueeze(1)

# tensor diagrams
t_d = td.to_diagram(t).fill_color(lambda idx, v: color_tensor[idx])
slice_1_d = td.to_diagram(slice_1).fill_color(lambda idx, v: color_tensor[idx])
slice_2_d = td.to_diagram(slice_2).fill_color(lambda idx, v: color_tensor[idx[0], 1, idx[2]])

# chalk diagrams
diagrams = []
for d in [t_d, slice_1_d, slice_2_d]:
  diagrams.append(d.to_chalk_diagram().center_xy())

# composite diagram
composite = chalk.hcat(diagrams, 1.0)

# add background + display
composite = composite.pad(1.5).center_xy()
env = composite.get_envelope()
chalk.set_svg_height(300)
chalk.rectangle(env.width * 0.8, env.height).fill_color(Color("white")) + composite.translate(dx=env.width * 0.1, dy=0)

For more examples and documentation, refer to this guide.

dev

pip install -e ".[dev]"

test

setup

Visual regression tests compare rendered output against reference images stored in fixtures/.

To generate or update reference images:

# all reference images
python tests/generate_references.py --all

# list reference images
python tests/generate_references.py --list

# select reference images
python tests/generate_references.py 3d_tensor styled_gradient

run

All tests:

pytest

license

MIT