← Software

TensorBloom

TensorBloom is a desktop PyTorch graph editor. I built its shape-analysis and code-generation paths to catch invalid models before training.

Personal project v0.1.0 2026

A graph editor has to keep four representations aligned: the visible graph, inferred tensor shapes, generated Python, and the PyTorch module used for training. TensorBloom recalculates shapes and regenerates code after every edit. The training module is built later from graph JSON, so that agreement is an invariant to maintain, not one the app verifies directly.

A graph is a tensor contract

Each edge asserts that the source tensor satisfies the destination operation's input requirements. TensorBloom topologically sorts the graph, propagates shapes from dataset or input nodes, and checks dimensions with operation-specific TypeScript rules.

Consider a deliberately small failure. MNIST supplies one-channel images, but the next convolution is configured to receive three channels. The graph is connected, yet the model cannot run as written.

Root error and deferred checks
StageContractInterpretation
Dataset[N, 1, 28, 28]One input channel
Conv2din_channels = 3Root mismatch: received 1
DownstreamShape depends on Conv2dDeferred until the root error is fixed

TensorBloom records the convolution as the root error and defers checks downstream by marking them as cascaded. This reduces duplicate-looking diagnostics, but it does not prove that every downstream node is valid or that every later error has the same cause. The property panel can propose the upstream channel count as a direct correction. Before training, a narrower preflight handles a small set of safe parameter fixes and stops if non-fixable errors remain.

One graph, several representations

The desktop application has three layers: a React and TypeScript graph editor, a small Rust and Tauri shell, and a Python and PyTorch sidecar connected through JSON-RPC. The TypeScript editor owns the live graph and shape rules. Its decompiler regenerates an nn.Module, including __init__ and forward, whenever nodes or edges change. For training, the sidecar constructs a PyTorch module from a handle-aware graph payload. Export rebuilds a module through a separate, narrower serialization path.

TensorBloom with a LeNet-5 graph above the generated PyTorch module and forward method
Figure 1. The graph and its generated PyTorch remain two views of the same model. This LeNet-5 project shows the graph above the regenerated module and forward method.TensorBloom v0.1.0 interface.Open full-size figure

Branches, merges, custom blocks, and multiple-output graphs affect name generation, data flow, shape propagation, source order, and the JSON-RPC payload sent to Python. These paths have to change together.

Where shape inference stops

TensorBloom infers shapes with TypeScript rules rather than executing arbitrary PyTorch. Unknown operations pass through their first input shape. Converted code blocks retain their original module type, so recognized types may continue to use the original rule even after their custom forward body changes. A displayed shape is therefore a prediction rather than proof of runtime validity. The interface should distinguish checked constraints from inferred values and unsupported cases. Otherwise, the graph overstates what the analyzer knows.

Current boundaries

  • The v0.1 training path supports static graphs, one loss node, and one GPU.
  • GAN and other multi-model training loops are outside the current abstraction.
  • Transformer attention-mask inputs are not represented in the training graph.
  • The ONNX and TorchScript export path omits code-block definitions and edge-handle metadata, does not apply the training path's dataset-to-input rewrite, and rebuilds the architecture with newly initialized weights.

CI builds the TypeScript frontend, checks the Rust shell, and exercises the supplied architecture templates with PyTorch forward passes. The project does not yet publish coverage measurements.

The hard part is consistency

The visual graph, inferred shapes, generated program, runtime module, and saved file are separate representations of one model. TensorBloom is reliable only when an edit has the same meaning in all five.

Failures expose that boundary better than successful graphs. A useful diagnostic identifies the first broken contract, separates its downstream effects, and states which parts of the graph were never checked.