a VS Code, Cursor & Antigravity extension

Your model,
drawn as it runs.

tensorViz is an extension for VS Code, Cursor and Antigravity that turns your PyTorch model into a live, shape-verified graph, right beside the file that defines it. Edit the graph, it patches your code. Save the file, the graph redraws from the tensors that actually flowed.

VS Code · Cursor · Antigravity · PyTorch · Runs on your machine
Conv2d 3→64[B, 64, 56, 56]
BatchNorm2d[B, 64, 56, 56]
ReLU[B, 64, 56, 56]
Linear →10~[B, 10]  (?)
measured ✓
run it to know →
σ(Wx + b)
scroll
see it · debug it

See the shape,
debug it in seconds.

PyTorch shapes fail silently, deep in the code. tensorViz makes the whole model visible, so the mismatch is obvious at a glance instead of three hours into a run.

01

Hallucinated shapes

Coding agents write layers instantly, then quietly mismatch a tensor dimension buried inside a nested block. You trace it by hand.

RuntimeError: mat1 and mat2 shapes cannot be multiplied
02

Wasted compute

A shape bug discovered mid-run burns GPU budget and kills momentum. The maths was wrong before the first epoch ever started.

3h 12m in · run halted · $$$
03

Blind to the structure

Code is a wall of text. The architecture, the thing you're actually reasoning about, stays invisible until it breaks.

where does this tensor even go?
bidirectional

The graph and the code
are the same thing.

Every node carries the file:line that produced it. Click a node, the cursor jumps there. Edit a parameter in the graph, it comes back as a minimal patch to your .py, one edit you can undo.

  • Click a node → cursor moves to its line
  • Move the cursor → the node highlights
  • Type a shape → your source updates
  • An edit that can't be made honestly is refused, with a reason
net.py ↔ graph synced
class Net(nn.Module):
    def __init__(self):
        super().__init__()
        self.conv = nn.Conv2d(3, 64, 3)   → 64×56×56
        self.bn   = nn.BatchNorm2d(64)  → 64×56×56
        self.fc   = nn.Linear(8, 8)      → 2×8
        self.act  = nn.ReLU()
    def forward(self, x):
        return self.act(self.fc(x))
saved · re-ran in your interpreter · redrawn from real tensors
the real thing

Your editor, watching your model.

Open a file, the graph appears beside it. Click a node, the code jumps. Edit the code, the graph follows. Here it is, driving itself.

tensorViz — models
Explorer
models
tiny_mlp.py
small_cnn.py
tests
README.md
tiny_mlp.py ×
▦ GRAPH · TinyMLP
⑃ main tiny_mlp.py shapes verified Python 3.12 tensorViz
what it does

An extension that tells
the truth about your model.

Measured shapes, not guesses

A save re-runs your class in a subprocess and redraws the graph from what actually flowed through it. Observed shapes appear as inlay hints on the lines that declared the layers. Only observed ones: an inferred shape in the same font would be a lie.

measured inferred

Import anything

Any installed torchvision, transformers, diffusers or timm model, or a cloned repo, becomes a shape-annotated graph. It reads the architecture from the code alone. Nothing is downloaded or run.

Verification pipeline

Shape inference plus a synthetic forward pass. Failing layers light up on the canvas with the actual torch traceback on the node.

Runnable export

Download a self-contained zip: model.py, sample_run.py, README.md, requirements.txt. It just runs.

Optional AI assist

Bring your own key for chat-driven edits over the same verifiable graph. Off by default, and every change it makes stays shape-checked.

what you get

Inside tensorViz.

Five surfaces you work in every day. Drag, or let it drift.

✳ Agent
Make the classifier head deeper.
✳ tensorViz I'll insert a Linear(128, 256) + ReLU before the output and re-verify the shapes.
+ Linear+ ReLU✓ shapes ok

AI agent

Chat that edits the verifiable graph, on your key.

Explorer
TENSORVIZ
▦ Open as Graph
class TinyMLP
▦ Open as Graph
class SmallCNN
▦ Open as Graph
class ResBlock

Lives in your editor

Every nn.Module gets an Open as Graph lens.

net.py
def __init__(self):
    self.conv = nn.Conv2d(3,64,3)  → 64×56×56
    self.bn   = nn.BatchNorm2d(64) → 64×56×56
    self.fc   = nn.Linear(8,10)   → 2×10

Measured shape hints

Real tensor shapes, inline on the lines that made them.

verify
Conv2d [B,64,56,56]
Linear shape error
RuntimeError mat1 and mat2 shapes cannot be multiplied (2x8 and 128x10)

Verification pipeline

Failing layers light up with the real torch traceback.

export
model.zip
model.py
sample_run.py
README.md
requirements.txt
$ python sample_run.py✓ runs

Runnable export

A self-contained zip that just runs.

free while in early access

Build neural networks
you can see.

Draw it. Verify the shapes. Export runnable PyTorch, all inside your editor.

[B, 3, 224, 224]
∑ wᵢxᵢ
→ verified
the extension is step one

tensorViz starts in your editor. It doesn't end there.

The extension is the first piece of a larger idea: making model architecture something you can see, verify, and reason about everywhere you work. More is on the way.

More to come