tensorViz: PyTorch model graphs with source navigation in VS Code

[tensorViz (tensorviz)]

Hi, I'm the maker of tensorViz, a free VS Code extension for reading PyTorch models.

When a model has layers nested inside layers, a flat printout makes it hard to connect the architecture to its Python code. tensorViz lets you expand one block, inspect its captured tensor shapes, and jump from a graph node to the corresponding source. Moving between the graph and editor is the main workflow.

Source-navigation demo

Runtime shapes come from a captured forward pass with example inputs. They describe that execution, not every possible branch or input. Source-only inspection is available separately; runtime capture needs a Python environment with PyTorch and the model's dependencies.

• Install from the VS Code Marketplace
• Explore example graphs in your browser

Where does a graph help you read unfamiliar PyTorch code, and where do you still need to trace it by hand? I'd welcome feedback on the graph-to-source workflow.

Posted with AI assistance for the tensorViz project.