Summary

TorchLens provides deep introspection into PyTorch models by capturing all activations and gradients during forward and backward passes. It automatically visualizes the computational graph, exposes extensive per-operation metadata, and allows for live interventions during model execution. This tool is invaluable for researchers and engineers needing to understand, debug, or modify complex, dynamic, or recurrent neural network architectures.

How It Works

TorchLens utilizes eager-mode Python-level function wrapping, intercepting operations as they execute. This approach is key to its ability to handle dynamic control flow, recurrent structures, and any PyTorch model that runs as standard Python code, differentiating it from static-graph analysis tools. It captures detailed information from both forward and backward passes, enabling granular analysis and real-time manipulation of model behavior.

Quick Start & Requirements

Installation requires Graphviz (e.g., sudo apt install graphviz) followed by pip install torchlens. The library is compatible with PyTorch 1.8.0+. Key resources include the official paper, a 10-minute tutorial notebook, facets tutorial, a gallery of visualizations, and detailed documentation on intervention APIs, backward capture, performance, and limitations.

Highlighted Details

• Comprehensive Data Capture: Records over 180 metadata fields per operation (550+ total), including shapes, dtypes, devices, timing, FLOPs, parameter info, and call stack context.
• Forward & Backward Pass Analysis: Captures gradients alongside activations using a unified API, enabling detailed gradient flow analysis through specific operations.
• Live Interventions: Allows dynamic modification of activations during the forward pass (e.g., zero-ablation, scaling, steering) via a rich set of helper functions and selectors.
• Automatic Graph Visualization: Generates PDF visualizations of the computational graph, supporting 'rolled' (compact) and 'unrolled' (detailed) modes for various architectures.
• Multi-Backend Support: Offers preview support for JAX, tinygrad, MLX, Paddle, and TensorFlow, with PyTorch as the primary, full-featured backend.

Maintenance & Community

TorchLens is stated to be in active development. Users are encouraged to report issues, ask questions, and provide suggestions via email, Twitter, the GitHub issues page, or the discussion board.

Licensing & Compatibility

The license type is not explicitly stated in the provided README. TorchLens is incompatible with torch.compile'd models, TorchScript, and torch.export. Specific behaviors and potential issues are noted for FSDP, sparse tensors, quantization, and torch.func.vmap.

Limitations & Caveats

The library does not support models processed by torch.compile, TorchScript, or torch.export due to their non-standard Python execution. Preview backends (JAX, tinygrad, etc.) may have fewer features than the PyTorch backend. Portable bundles contain pickle files and pose a security risk if loaded from untrusted sources.