### Install Torchvista via Pip

Source: https://github.com/sachinhosmani/torchvista/blob/main/README.md

Install the torchvista package using pip. This is the standard method for installing Python packages.

```bash
pip install torchvista
```

--------------------------------

### Basic Tensor Visualization

Source: https://github.com/sachinhosmani/torchvista/blob/main/docs/templates/tutorial.html

This snippet shows how to visualize a basic PyTorch tensor using TorchVista. Ensure TorchVista is installed and imported.

```python
import torch
from torchvista import TorchVista

# Create a sample tensor
tensor = torch.randn(3, 224, 224)

# Initialize TorchVista
tv = TorchVista()

# Visualize the tensor
tv.visualize(tensor)
```

--------------------------------

### Install Torchvista via Conda

Source: https://github.com/sachinhosmani/torchvista/blob/main/README.md

Install the torchvista package using conda from the conda-forge channel. Refer to the feedstock for more details.

```bash
conda install -c conda-forge torchvista
```

--------------------------------

### Trace a SimpleResNet Model

Source: https://github.com/sachinhosmani/torchvista/blob/main/docs/generated/tutorials/exploring_modules.html

Instantiates a SimpleResNet model and an example input, then uses torchvista.trace_model to visualize the model structure.

```python
model = SimpleResNet()
example_input = torch.randn(1, 3, 32, 32)

trace_model(model, example_input)
```

--------------------------------

### trace_model - Basic Usage

Source: https://context7.com/sachinhosmani/torchvista/llms.txt

The primary entry point for torchvista. Runs the model's forward pass with the provided inputs, builds an interactive computation graph, and renders it inline in the current notebook cell. This example demonstrates basic usage with a simple two-branch model.

```APIDOC
## trace_model - Visualize a PyTorch model's forward pass

### Description
The primary entry point for torchvista. Runs the model's forward pass with the provided inputs, builds an interactive computation graph, and renders it inline in the current notebook cell. All parameters beyond `model` and `inputs` are optional and control the visual appearance, tracing depth, and export behavior.

### Method
`trace_model(model, inputs, ...)`

### Parameters
- **model** (torch.nn.Module) - The PyTorch model to trace.
- **inputs** (torch.Tensor or tuple or dict) - The input(s) to the model's forward method.

### Request Example
```python
from torchvista import trace_model
import torch
import torch.nn as nn

class LinearModel(nn.Module):
    def __init__(self):
        super().__init__()
        self.linear1 = nn.Linear(10, 5)
        self.linear2 = nn.Linear(10, 5)

    def forward(self, x):
        return self.linear1(x) + self.linear2(x)

model = LinearModel()
inputs = torch.randn(2, 10)

# Renders an interactive graph inline in the notebook
trace_model(model, inputs)
```
```

--------------------------------

### Control Initial Visual Expansion with collapse_modules_after_depth

Source: https://context7.com/sachinhosmani/torchvista/llms.txt

Controls how deeply nested modules are initially expanded in the rendered graph. Modules nested beyond collapse_modules_after_depth levels start collapsed but can be expanded interactively. Setting this equal to forced_module_tracing_depth expands all traced modules immediately.

```python
import torch
import torch.nn as nn
from torchvista import trace_model

class ResidualBlock(nn.Module):
    def __init__(self, ch):
        super().__init__()
        self.conv = nn.Sequential(
            nn.Conv2d(ch, ch, 3, padding=1),
            nn.BatchNorm2d(ch),
            nn.ReLU(),
            nn.Conv2d(ch, ch, 3, padding=1),
            nn.BatchNorm2d(ch),
        )
        self.relu = nn.ReLU()

    def forward(self, x):
        return self.relu(x + self.conv(x))

class SimpleResNet(nn.Module):
    def __init__(self):
        super().__init__()
        self.stem   = nn.Sequential(nn.Conv2d(3, 16, 3, padding=1))
        self.block1 = ResidualBlock(16)
        self.block2 = ResidualBlock(16)
        self.pool   = nn.AdaptiveAvgPool2d((1, 1))
        self.fc     = nn.Linear(16, 10)

    def forward(self, x):
        x = self.stem(x)
        x = self.block1(x)
        x = self.block2(x)
        x = self.pool(x).squeeze(-1).squeeze(-1)
        return self.fc(x)

model  = SimpleResNet()
inputs = torch.randn(1, 3, 32, 32)

# Trace to depth 3 AND pre-expand the visualization to depth 3
trace_model(
    model,
    inputs,
    forced_module_tracing_depth=3,      # trace inside Conv2d, BN, etc.
    collapse_modules_after_depth=3,     # show all levels expanded by default
)
# collapse_modules_after_depth=0 collapses everything; =1 (default) expands only top level
```

--------------------------------

### Trace a PyTorch Model for Visualization

Source: https://github.com/sachinhosmani/torchvista/blob/main/docs/index.html

Trace a defined PyTorch model with an example input to generate a visualization. This function is the core of TorchVista's visualization capability.

```python
example_input = torch.randn(1, 10, 32)
trace_model(model, example_input)
```

--------------------------------

### trace_model - Tuple Inputs

Source: https://context7.com/sachinhosmani/torchvista/llms.txt

When a model's `forward` method accepts multiple positional arguments, pass inputs as a Python tuple. Torchvista unpacks the tuple and calls `model(*inputs)`. This example shows how to handle models with multiple input arguments.

```APIDOC
## trace_model - Tuple inputs for multi-argument models

### Description
When a model's `forward` method accepts multiple positional arguments, pass inputs as a Python tuple. Torchvista unpacks the tuple and calls `model(*inputs)`.

### Method
`trace_model(model, inputs, ...)`

### Parameters
- **model** (torch.nn.Module) - The PyTorch model to trace.
- **inputs** (tuple) - A tuple containing the multiple inputs for the model's forward method.

### Request Example
```python
import torch
import torch.nn as nn
from torchvista import trace_model

class TwoInputModel(nn.Module):
    def __init__(self):
        super().__init__()
        self.fc = nn.Linear(20, 5)

    def forward(self, x, y):
        return self.fc(torch.cat([x, y], dim=-1))

model = TwoInputModel()
x = torch.randn(2, 10)
y = torch.randn(2, 10)

# Pass multiple inputs as a tuple
trace_model(model, (x, y))
```
```

--------------------------------

### Define a PyTorch Model with Self-Attention

Source: https://github.com/sachinhosmani/torchvista/blob/main/docs/index.html

Define a custom PyTorch model, including a SelfAttention module and an AttentionClassifier. This serves as an example model to be traced and visualized.

```python
import torch
import torch.nn as nn
from torchvista import trace_model

class SelfAttention(nn.Module):
    def __init__(self, embed_dim):
        super().__init__()
        self.query = nn.Linear(embed_dim, embed_dim)
        self.key = nn.Linear(embed_dim, embed_dim)
        self.value = nn.Linear(embed_dim, embed_dim)
        self.scale = embed_dim ** 0.5

    def forward(self, x):
        q = self.query(x)
        k = self.key(x)
        v = self.value(x)
        attn = torch.softmax((q @ k.transpose(-2, -1)) / self.scale, dim=-1)
        return attn @ v

class AttentionClassifier(nn.Module):
    def __init__(self):
        super().__init__()
        self.embed = nn.Linear(32, 64)
        self.attn = SelfAttention(64)
        self.classifier = nn.Linear(64, 5)

    def forward(self, x):
        x = self.embed(x)
        x = self.attn(x)
        return self.classifier(x.mean(dim=1))
```

--------------------------------

### Basic Model Tracing with Torchvista

Source: https://github.com/sachinhosmani/torchvista/blob/main/README.md

Import torchvista and use the trace_model function to visualize a PyTorch model's forward pass. Requires PyTorch and a defined model with input tensors.

```python
import torch
import torch.nn as nn

# Import torchvista
from torchvista import trace_model

# Define your module
class LinearModel(nn.Module):
    def __init__(self):
        super().__init__()
        self.linear = nn.Linear(10, 5)

    def forward(self, x):
        return self.linear(x)

# Instantiate the module and tensor input
model = LinearModel()
inputs = torch.randn(2, 10)

# Trace!
trace_model(model, inputs)
```

--------------------------------

### Initialize Graph Visualization Variables

Source: https://github.com/sachinhosmani/torchvista/blob/main/torchvista/templates/graph.html

Sets up global variables and constants for graph rendering, including data mappings, color schemes, and configuration flags. It also initializes sets and maps for tracking node states and container hierarchies.

```javascript
}; const graph_node_name_to_without_suffix = ${graph_node_name_to_without_suffix}; const graph_node_display_names = ${graph_node_display_names}; const node_to_attr_name = ${node_to_attr_name}; const ancestor_map = ${ancestor_map}; const repeat_containers = new Set(${repeat_containers}); const REPEAT_CONTAINER_COLOR = "#FA8334"; const collapse_modules_after_depth = ${collapse_modules_after_depth}; const show_module_attr_names = ${show_module_attr_names}; const node_to_module_path = ${node_to_module_path}; const generateImage = ${generate_image}; const generateSvg = ${generate_svg}; const show_modular_view = ${show_modular_view}; var collapsedContainers = new Set(); const containerNodes = new Set(); var adj_list_collapsed_nodes = {}; let containerHierarchy = {}; let vizInstance = null; let containerNodeCounts = {}; const EXPAND_TARGET_COVERAGE = 0.8; const COLLAPSE_TARGET_COVERAGE = 0.2; let lastModifiedContainer = null; // Track the last modified container let lastOperation = null; // Track whether it was 'expand' or 'collapse'
```

--------------------------------

### Basic Usage of trace_model

Source: https://context7.com/sachinhosmani/torchvista/llms.txt

Demonstrates the fundamental use of trace_model with a simple PyTorch model and tensor inputs. Ensure torchvista, torch, and torch.nn are imported.

```python
from torchvista import trace_model
import torch
import torch.nn as nn

# --- Basic usage: simple two-branch model ---
class LinearModel(nn.Module):
    def __init__(self):
        super().__init__()
        self.linear1 = nn.Linear(10, 5)
        self.linear2 = nn.Linear(10, 5)

    def forward(self, x):
        return self.linear1(x) + self.linear2(x)

model = LinearModel()
inputs = torch.randn(2, 10)

# Renders an interactive graph inline in the notebook
trace_model(model, inputs)
```

--------------------------------

### Tuple Inputs for Multi-Argument Models with trace_model

Source: https://context7.com/sachinhosmani/torchvista/llms.txt

Illustrates how to pass multiple positional arguments to a model's forward method by providing inputs as a Python tuple. Torchvista will unpack this tuple for the model call.

```python
import torch
import torch.nn as nn
from torchvista import trace_model

class TwoInputModel(nn.Module):
    def __init__(self):
        super().__init__()
        self.fc = nn.Linear(20, 5)

    def forward(self, x, y):
        return self.fc(torch.cat([x, y], dim=-1))

model = TwoInputModel()
x = torch.randn(2, 10)
y = torch.randn(2, 10)

# Pass multiple inputs as a tuple
trace_model(model, (x, y))
```

--------------------------------

### Resize Observer for JSON Viewer

Source: https://github.com/sachinhosmani/torchvista/blob/main/torchvista/templates/graph.html

Sets up a ResizeObserver to monitor the popup element. When the popup resizes, it iterates through any JSON viewers within it and calls their resize method to ensure proper rendering.

```javascript
const popup = document.getElementById(
`popup_${unique_id}`
);
const resizeObserver = new ResizeObserver(() => {
  const editorContainers = popup.querySelectorAll('.json-viewer_${unique_id}');
  for (let container of editorContainers) {
    if (container.jsonEditor) {
      container.jsonEditor.resize();
    }
  }
});
resizeObserver.observe(popup);
```

--------------------------------

### Show Popup from Info Button

Source: https://github.com/sachinhosmani/torchvista/blob/main/torchvista/templates/graph.html

Displays a popup window for a selected node, centering the node in the view and highlighting it. Populates the popup with node information, including module path and attribute name.

```javascript
function showPopupFromInfoButton(event, [nodeName], infoButtonElement = null) {
  event.stopPropagation();
  const clickedNodeElement = d3.select(this);
  // Position based on info button or node center
  const clickednodeData = graphvizData.nodes[nodeName];
  if (clickednodeData) {
    const svgNode = svg.node();
    const width = svgNode.clientWidth || 800;
    const height = svgNode.clientHeight || 600;
    const currentTransform = d3.zoomTransform(svg.node());
    let targetX, targetY;
    if (infoButtonElement) {
      // Calculate info button position relative to node center
      const infoButtonOffsetX = clickednodeData.width / 2 - 8;
      const infoButtonOffsetY = -clickednodeData.height / 2 + 8;
      // Calculate where the info button should be positioned on screen
      const infoButtonWorldX = clickednodeData.position.x + infoButtonOffsetX;
      const infoButtonWorldY = clickednodeData.position.y + infoButtonOffsetY;
      targetX = 0.4 * width - infoButtonWorldX * currentTransform.k;
      targetY = 0.2 * height - infoButtonWorldY * currentTransform.k;
    } else {
      // Original behavior - center the node
      targetX = 0.4 * width - clickednodeData.position.x * currentTransform.k;
      targetY = 0.2 * height - clickednodeData.position.y * currentTransform.k;
    }
    svg.call(zoom.transform, d3.zoomIdentity
      .translate(targetX, targetY)
      .scale(currentTransform.k));
  }

  clickedNodeElement.classed("highlight-node_${unique_id}", true);
  setTimeout(() => {
    clickedNodeElement.classed("highlight-node_${unique_id}", false);
  }, 2000);

  const popup = d3.select("#popup_${unique_id}");
  const moduleTemplate = d3.select("#module-template_${unique_id}");
  const tensorOpTemplate = d3.select("#tensor-op-template_${unique_id}");
  const noInfoMessage = d3.select("#no-info_${unique_id}");

  const nodeData = adj_list[nodeName];
  let popupTitle = nodeName;
  if (node_to_module_path[nodeName] && graph_node_name_to_without_suffix[nodeName]) {
    popupTitle = node_to_module_path[nodeName] + "." + graph_node_name_to_without_suffix[nodeName];
  }
  d3.select("#popup-title_${unique_id}").text(popupTitle);

  // Show attribute name if available
  const attrName = node_to_attr_name[nodeName];
  const attrNameRow = d3.select("#attr-name-row_${unique_id}");
  if (attrName) {
    d3.select("#attr-name-value_${unique_id}").text(attrName);
    attrNameRow.style("display", "block");
  } else {
    attrNameRow.style("display", "none");
  }

  moduleTemplate.style("display", "none");
  tensorOpTemplate.style("display", "none");
  noInfoMessage.style("display", "none");

  if (module_info[nodeName]) {
    const info = 
```

--------------------------------

### Define and Trace Model with Custom Collapse Depth

Source: https://github.com/sachinhosmani/torchvista/blob/main/docs/generated/tutorials/collapse_modules_after_depth.html

Define a simple neural network and trace it using `torchvista.trace_model`. Set `collapse_modules_after_depth` to control which modules are initially collapsed in the visualization. This is useful for managing the complexity of large models.

```python
import torch
import torch.nn as nn
from torchvista import trace_model

class ResidualBlock(nn.Module):
    def __init__(self, channels):
        super().__init__()
        self.conv = nn.Sequential(
            nn.Conv2d(channels, channels, 3, padding=1),
            nn.BatchNorm2d(channels),
            nn.ReLU(),
            nn.Conv2d(channels, channels, 3, padding=1),
            nn.BatchNorm2d(channels),
        )
        self.relu = nn.ReLU()

    def forward(self, x):
        return self.relu(x + self.conv(x))

class SimpleResNet(nn.Module):
    def __init__(self):
        super().__init__()
        self.stem = nn.Sequential(
            nn.Conv2d(3, 16, 3, padding=1)
        )
        self.block1 = ResidualBlock(16)
        self.block2 = ResidualBlock(16)
        self.pool = nn.AdaptiveAvgPool2d((1, 1))
        self.fc = nn.Linear(16, 10)

    def forward(self, x):
        x = self.stem(x)
        x = self.block1(x)
        x = self.block2(x)
        x = self.pool(x).squeeze(-1).squeeze(-1)
        return self.fc(x)

model = SimpleResNet()
example_input = torch.randn(1, 3, 32, 32)

# Expand modules up to depth 3 when first displayed
trace_model(
    model,
    example_input,
    forced_module_tracing_depth=3,
    ##############################
    collapse_modules_after_depth=3  # <-- modules beyond this depth start collapsed
    ##############################
)
```

--------------------------------

### Compress Repeated Layer Patterns with `show_compressed_view`

Source: https://context7.com/sachinhosmani/torchvista/llms.txt

Use `show_compressed_view=True` to collapse consecutive identical layers within `nn.Sequential` or `nn.ModuleList` into `REPEAT Nx` blocks. This feature is experimental and may impact performance on very large models.

```python
import torch
import torch.nn as nn
from torchvista import trace_model

class DeepModel(nn.Module):
    def __init__(self):
        super().__init__()
        # 10 identical Sequential blocks, each with 10 Linear layers
        block = nn.Sequential(*[nn.Linear(64, 64) for _ in range(10)])
        self.layers = nn.ModuleList([block] * 10)

    def forward(self, x):
        for layer in self.layers:
            x = layer(x)
        return x

model  = DeepModel()
inputs = torch.randn(2, 64)

# Without compression: 100 Linear nodes visible
trace_model(model, inputs)

# With compression: repeated patterns collapsed into "REPEAT 10x" blocks
trace_model(
    model,
    inputs,
    show_compressed_view=True,   # collapses repeated Sequential/ModuleList patterns
)
# Note: only works for chains within nn.Sequential or nn.ModuleList
```

--------------------------------

### Trace a Simple PyTorch Model with Torchvista

Source: https://github.com/sachinhosmani/torchvista/blob/main/docs/generated/tutorials/basic_usage.html

Use `trace_model` to visualize the forward pass of a PyTorch model. This is useful for understanding model architecture and data flow.

```python
import torch
import torch.nn as nn
from torchvista import trace_model

class LinearModel(nn.Module):
    def __init__(self):
        super().__init__()
        self.linear1 = nn.Linear(10, 5)
        self.linear2 = nn.Linear(10, 5)

    def forward(self, x):
        return self.linear1(x) + self.linear2(x)

model = LinearModel()
example_input = torch.randn(2, 10)

# Visualize the forward pass
trace_model(model, example_input)
```

--------------------------------

### Graph Data Initialization

Source: https://github.com/sachinhosmani/torchvista/blob/main/torchvista/templates/graph.html

Initializes graph rendering by defining adjacency list, module information, function information, and parent module mappings. These variables are populated from JSON data.

```javascript
(function() {
const adj_list = ${adj_list_json};
const module_info = ${module_info_json};
const func_info = ${func_info_json};
const parent_module_to_nodes = ${parent_module_to_nodes_json};
const parent_module_to_depth = ${parent_module_to_depth_json}
```

--------------------------------

### Create and Render Graph with D3.js

Source: https://github.com/sachinhosmani/torchvista/blob/main/torchvista/templates/graph.html

Initializes an SVG canvas, sets up zoom behavior, defines an arrowhead marker, and renders graph clusters and edges using D3.js. This is the main function for graph visualization.

```javascript
function createGraph(graphvizData) {
 const svg = d3.select("#graph_${unique_id}")
 .append("svg")
 .attr("width", "100%")
 .attr("height", "100%")
 .style("background-color", "#f3f3eb")
 .style("height", "100%");

 const inner = svg.append("g");

 // Add zoom behavior
 const zoom = d3.zoom().on("zoom", function(event) {
 removeTooltip();
 inner.attr("transform", event.transform);
 closePopup();
 });
 svg.call(zoom);

 // Define arrowhead marker
 svg.append("defs").append("marker")
 .attr("id", "arrowhead_${unique_id}")
 .attr("viewBox", "0 -5 10 10")
 .attr("refX", 8)
 .attr("refY", 0)
 .attr("markerWidth", 6)
 .attr("markerHeight", 6)
 .attr("orient", "auto")
 .append("path")
 .attr("d", "M0,-5L10,0L0,5")
 .attr("fill", "#333");

 // Draw clusters (container boxes) first so they're behind nodes
 const clusters = inner.selectAll(".boundary-box_${unique_id}")
 .data(Object.entries(graphvizData.clusters))
 .enter()
 .append("g")
 .attr("class", "boundary-box_${unique_id}");

 const containersWithFailedNodes = identifyContainersWithFailedNodes(adj_list, ancestor_map);

 clusters.each(function([clusterName, clusterData]) {
 const group = d3.select(this);
 const hasFailedNodes = containersWithFailedNodes.has(clusterName);
 const isRepeat = isRepeatNode(clusterName);

 if (hasFailedNodes) {
 group.classed("error_${unique_id}", true);
 }
 if (isRepeat) {
 group.classed("repeat-container_${unique_id}", true);
 }

 // Parse bounding box coordinates
 const [x1, y1, x2, y2] = clusterData.bb.split(",").map(Number);

 // Create boundary box
 const rect = group.append("rect")
 .attr("x", x1)
 .attr("y", y1)
 .attr("width", x2 - x1)
 .attr("height", y2 - y1)
 .attr("fill-opacity", 0.0)
 .attr("stroke", isRepeat ? REPEAT_CONTAINER_COLOR : "black")
 .attr("stroke-dasharray", isRepeat ? "3 2" : null)
 .attr("stroke-width", isRepeat ? 1.2 : 1)
 .attr("rx", 5)
 .attr("ry", 5);

 if (isRepeat) {
 rect.style("filter", `drop-shadow(0 0 2px ${REPEAT_CONTAINER_COLOR})`);
 }

 // Create label for the cluster
 const textGroup = group.append("g")
 .attr("transform", `translate(${x1 + 10}, ${y1 + 3})
`);

 const text = textGroup.append("text")
 .text(formatContainerDisplay(clusterData.label))
 .attr("text-anchor", "start")
 .attr("dominant-baseline", "hanging")
 .attr("dy", "-6")
 .attr("font-size", 7)
 .attr("fill", "#333")
 .attr("stroke", "#f3f3eb")
 .attr("stroke-width", "8px")
 .attr("paint-order", "stroke")
 .attr("stroke-linejoin", "round");

 if (!isRepeatNode(clusterName)) {
 // Add collapse/expand button
 const collapseButton = group.append("circle")
 .attr("class", "collapse-button_${unique_id}")
 .attr("cx", x1 + 10)
 .attr("cy", y1 + 10)
 .attr("r", 5.5);

 group.append("text")
 .attr("class", "collapse-icon_${unique_id}")
 .attr("x", x1 + 10)
 .attr("y", y1 + 10)
 .attr("text-anchor", "middle")
 .attr("dominant-baseline", "middle")
 .attr("font-size", "10px")
 .text("-");

 // Add custom tooltip events
 collapseButton.on("mouseenter", function() {
 showTooltip(this, "Collapse module");
 })
 .on("mouseleave", removeTooltip);
 }
 });

 // Add click event handler for collapse buttons
 clusters.selectAll(".collapse-button_${unique_id}, .collapse-icon_${unique_id}")
 .on("click", function(event, [clusterName, clusterData]) {
 event.stopPropagation(); // Prevent event from bubbling to SVG
 removeTooltip();
 if (isRepeatNode(clusterName)) return;

 // Collapse the container
 collapsedContainers.add(clusterName);
 // Track this container and operation for centering after reload
 lastModifiedContainer = clusterName;
 lastOperation = "collapse";
 reloadGraph();
 });

 // Create edges using the path points from Graphviz
 const edges = inner.selectAll(".edge_${unique_id}")
 .data(graphvizData.edges)
 .enter()
 .append("g")
 .attr("class", "edge-group_${unique_id}");

 // Add the actual path
 edges.append("path")
 .attr("class", function(edge) {
 if (edge.is_implied_edge) {
 return "edge_${unique_id} implied-edge_${unique_id}";
 }
 return "edge_${unique_id}";
 })
 .attr("d", function(edge) {
 if (edge.path && edge.path.length > 0) {
 // Create a proper SVG path using Bézier curves
 let d = "";
 const points = edge.path;
 // Start at the first point
 d += `M${points[0].x},${points[0].y}`;
 // Create cubic Bézier curves
 for (let i = 1; i < points.length - 2; i += 3) {
 if (i + 2 < points.length) {
 d += ` C${points[i].x},${points[i].y} ${points[i+1].x},${points[i+1].y} ${points[i+2].x},${points[i+2].y}`;
 }
 }
 // If we have leftover points, draw a line to the last point
 if (points.length % 3 !== 1) {
 d += ` L${points[points.length-1].x},${points[points.length-1].y}`;
 }
 return d;
 }
 return "";
 })
 .attr("marker-end", "url(#arrowhead_${unique_id})");

 // Add dimension labels on edges
 edges.each(function(edge) {
 const edgeGroup = d3.select(this);
 const path = edgeGroup.select("path").node();
 if (!path) return;

 // Find the source node data
 if (!adj_list[edge.source] && !adj_list_collapsed_nodes[edge.source]) return;

 // Find the edge data with matching target
 let sourceNodeData = adj_list_collapsed_nodes[edge.sou
```

--------------------------------

### Initialize and Reload Graph

Source: https://github.com/sachinhosmani/torchvista/blob/main/torchvista/templates/graph.html

Initializes graph rendering logic, including handling collapsed containers and choosing the appropriate rendering method (image, SVG, or interactive).

```javascript
function reloadGraph() {
  d3.select("#graph_${unique_id}").html("");
  if (generateImage) {
    renderGraphvizImage();
  } else if (generateSvg) {
    renderGraphvizSvg();
  } else {
    const dotSource = generateDotSource(adj_list, ancestor_map);
    graphData = vizInstance.renderJSON(dotSource) // Extract the graph data with positions from Graphviz
    const extractedGraphData = extractGraphData(graphData);
    createGraph(extractedGraphData); // Needed to get the right vp coverage when containers are expanded
    calculateContainerNodeCounts();
  }
}
```

```javascript
waitForLibs(function() {
  initContainerNodes();
  initializeCollapsedContainersFromGraph();
  reloadGraph();
});
```

--------------------------------

### Render Graphviz Image (PNG)

Source: https://github.com/sachinhosmani/torchvista/blob/main/torchvista/templates/graph.html

Generates and displays a PNG image of the graph. Includes error handling for large models and provides download/view links.

```javascript
function renderGraphvizImage() {
  d3.select("#png-link-container_${unique_id}").remove();
  function showErrorMessage(error) {
    const parent = d3.select("#graph-container_${unique_id}").node().parentNode;
    const errorContainer = d3.select(parent)
      .insert("div", "#graph-container_${unique_id}")
      .attr("id", "png-error-container_${unique_id}");
    errorContainer.append("div")
      .attr("id", "png-error-message_${unique_id}")
      .html(`<strong>Error:</strong><br>${error.stack || error.message || error}<br><br><strong>Note:</strong> Exporting very large models as PNG is currently not supported. Please use generate_image=False, or raise an issue on https://github.com/sachinhosmani/torchvista/issues`);
  }
  async function generatePNG() {
    processCollapsedGraph();
    const tbDotSource = generateDotFromProcessedData('TB');
    const svgString = vizInstance.renderString(tbDotSource, { format: 'svg' });
    const svgBlob = new Blob([svgString], { type: 'image/svg+xml' });
    const svgUrl = URL.createObjectURL(svgBlob);
    const img = new Image();
    const imageLoaded = new Promise((resolve, reject) => {
      img.onload = resolve;
      img.onerror = reject;
      img.src = svgUrl;
    });
    await imageLoaded;
    const canvas = document.createElement('canvas');
    const ctx = canvas.getContext('2d');
    const padding = 40;
    canvas.width = img.naturalWidth + (padding * 2);
    canvas.height = img.naturalHeight + (padding * 2);
    ctx.fillStyle = '#f3f3eb';
    ctx.fillRect(0, 0, canvas.width, canvas.height);
    ctx.drawImage(img, padding, padding);
    const pngBlob = await new Promise((resolve, reject) => {
      canvas.toBlob(resolve, 'image/png');
    });
    const pngUrl = URL.createObjectURL(pngBlob);
    // Create link container
    const parent = d3.select("#graph-container_${unique_id}").node().parentNode;
    const linkContainer = d3.select(parent)
      .insert("div", "#graph-container_${unique_id}")
      .attr("id", "png-link-container_${unique_id}");
    linkContainer.append("a")
      .attr("href", pngUrl)
      .attr("target", "_blank")
      .attr("id", "png-download-button_${unique_id}")
      .text("View PNG in New Tab");
    linkContainer.append("a")
      .attr("href", pngUrl)
      .attr("download", `torchvista_graph_${unique_id}.png`)
      .attr("id", "png-save-button_${unique_id}")
      .text("Download PNG");
    URL.revokeObjectURL(svgUrl);
  }
  generatePNG().catch(error => {
    console.error("Error generating PNG:", error);
    showErrorMessage(error);
  });
}
```

--------------------------------

### Define and Trace a Simple ResNet Model

Source: https://github.com/sachinhosmani/torchvista/blob/main/docs/generated/tutorials/show_non_gradient_nodes.html

Defines a simple ResNet model and traces it using `trace_model`. Set `show_non_gradient_nodes=False` to hide constant tensors and scalars that do not require gradients.

```python
import torch
import torch.nn as nn
from torchvista import trace_model

class ResidualBlock(nn.Module):
    def __init__(self, channels):
        super().__init__()
        self.conv = nn.Sequential(
            nn.Conv2d(channels, channels, 3, padding=1),
            nn.BatchNorm2d(channels),
            nn.ReLU(),
            nn.Conv2d(channels, channels, 3, padding=1),
            nn.BatchNorm2d(channels),
        )
        self.relu = nn.ReLU()

    def forward(self, x):
        return self.relu(x + self.conv(x))

class SimpleResNet(nn.Module):
    def __init__(self):
        super().__init__()
        self.stem = nn.Sequential(
            nn.Conv2d(3, 16, 3, padding=1)
        )
        self.block1 = ResidualBlock(16)
        self.block2 = ResidualBlock(16)
        self.pool = nn.AdaptiveAvgPool2d((1, 1))
        self.fc = nn.Linear(16, 10)

    def forward(self, inputs):
        x = inputs['image']
        x = self.stem(x)
        x = self.block1(x)
        x = self.block2(x)
        x = self.pool(x).squeeze(-1).squeeze(-1)
        return self.fc(x)

model = SimpleResNet()
example_input = {'image': torch.randn(1, 3, 32, 32)}

# Hide constant tensors and scalars that don't require gradients
trace_model(
    model,
    example_input,
    forced_module_tracing_depth=3,
    collapse_modules_after_depth=3,
    ##############################
    show_non_gradient_nodes=False  # <-- hides non-gradient nodes
    ##############################
)
```

--------------------------------

### Create JSON Viewer Component

Source: https://github.com/sachinhosmani/torchvista/blob/main/torchvista/templates/graph.html

Initializes a JSON editor instance within a given container element. Sets editor options for code mode, disabling menus, and setting indentation.

```javascript
function createJsonViewer(container, data) {
  // Clear previous contents
  container.innerHTML = '';
  // Create a new div to hold the editor
  const editorContainer = document.createElement('div');
  editorContainer.style.height = '180px';
  container.appendChild(editorContainer);

  const options = {
    mode: 'code',
    mainMenuBar: false,
    navigationBar: false,
    statusBar: false,
    indentation: 2,
    onChange: function() {
      editor.set(data);
    }
  };
  const editor = new JSONEditor(editorContainer, options);
  editor.set(data);
  // set a reference for later use during resizing
  container.jsonEditor = editor;
}
```

--------------------------------

### Enable Compressed View for Repeated Structures

Source: https://github.com/sachinhosmani/torchvista/blob/main/docs/generated/tutorials/compressed_view.html

Set `show_compressed_view=True` when calling `trace_model` to compress identical, sequential blocks within your model. This is particularly useful for models with many repeating layers.

```python
import torch
import torch.nn as nn
from torchvista import trace_model

class DeepModel(nn.Module):
    def __init__(self):
        super().__init__()
        # 10 identical Sequential blocks, each with 10 Linear layers
        block = nn.Sequential(*[nn.Linear(64, 64) for _ in range(10)])
        self.layers = nn.ModuleList([block] * 10)

    def forward(self, x):
        for seq in self.layers:
            x = seq(x)
        return x

model = DeepModel()
example_input = torch.randn(2, 64)

# Compress repeated structures into a single representation
trace_model(
    model,
    example_input,
    ###############################
    show_compressed_view=True  # <-- compresses repeated layers
    ###############################
)
```

--------------------------------

### trace_model

Source: https://github.com/sachinhosmani/torchvista/blob/main/README.md

Visualizes the forward pass of a PyTorch model directly in a web-based notebook. It allows for interactive exploration of the model's structure, parameters, and attributes, with options for customization and export.

```APIDOC
## trace_model

### Description
Visualizes the forward pass of a PyTorch model directly in a web-based notebook. It allows for interactive exploration of the model's structure, parameters, and attributes, with options for customization and export.

### Method
`trace_model`

### Parameters
#### Tracing Parameters
- **model** (torch.nn.Module) - Required - Model instance to visualize.
- **inputs** (Any) - Required - Input(s) forwarded into the model; pass a single input or a tuple.
- **forced_module_tracing_depth** (int | None) - Optional - Maximum depth of module internals to trace; `None` traces only user-defined modules.

#### Visual Parameters
- **show_non_gradient_nodes** (bool) - Optional - Display nodes for constants and other values outside the gradient graph. Defaults to `True`.
- **collapse_modules_after_depth** (int) - Optional - Depth to initially expand nested modules; `0` collapses everything (nodes can still be expanded interactively). Defaults to `1`.
- **height** (int) - Optional - Canvas height in pixels. Defaults to `800`.
- **width** (int | str | None) - Optional - Canvas width; accepts pixels or percentages; defaults to full available width when omitted.
- **show_module_attr_names** (bool) - Optional - Display attribute names for modules when available instead of just class names. Defaults to `False`.
- **show_compressed_view** (bool) - Optional (Experimental) - Compress the graph by showing repeating nodes of the same type with identical input and output dims in single "repeat" blocks. This feature currently only recognises repeating nodes within `Sequential` and `ModuleList`. WARNING: this feature might be expensive on large models. Defaults to `False`.

#### Export Parameters
- **export_format** (str | None) - Optional - Optional export format: `png`, `svg`, or `html` if exporting graph as a file. Otherwise, by default the graph is shown within the notebook. Defaults to `None`.
- **export_path** (str | None) - Optional - Custom path if exporting as a file. **Only HTML format** is currently supported with custom export paths. If only file name is specified, it will be created inside the present working directory. Defaults to `None`.

### Request Example
```python
import torch
import torch.nn as nn
from torchvista import trace_model

class LinearModel(nn.Module):
    def __init__(self):
        super().__init__()
        self.linear = nn.Linear(10, 5)

    def forward(self, x):
        return self.linear(x)

model = LinearModel()
inputs = torch.randn(2, 10)

trace_model(model, inputs, show_non_gradient_nodes=True, collapse_modules_after_depth=1)
```

### Response
This function does not return a value; it displays the visualization directly in the notebook or exports it based on the provided parameters.
```

--------------------------------

### Dict Inputs for Named Tensor Nodes with trace_model

Source: https://context7.com/sachinhosmani/torchvista/llms.txt

Shows how to use dictionary inputs to label tensor nodes in the visualization with custom names, improving clarity for complex models. Nested dictionaries create path-style labels.

```python
import torch
import torch.nn as nn
from torchvista import trace_model

class MultiModalModel(nn.Module):
    def __init__(self):
        super().__init__()
        self.image_encoder = nn.Linear(64, 32)
        self.text_encoder  = nn.Linear(32, 32)
        self.fusion        = nn.Linear(64, 16)

    def forward(self, inputs):
        img = inputs['visual']['image']        # node label: 'visual.image'
        txt = inputs['text']['embedding']      # node label: 'text.embedding'
        img_feat = self.image_encoder(img)
        txt_feat = self.text_encoder(txt)
        fused = torch.cat([img_feat, txt_feat], dim=-1)
        out = self.fusion(fused)
        # Dict output also produces named output nodes
        return {'results': {'prediction': out, 'features': [img_feat, txt_feat]}}

model = MultiModalModel()
example_input = {
    'visual': {'image':     torch.randn(1, 64)},
    'text':   {'embedding': torch.randn(1, 32)},
}

trace_model(model, example_input)
# Input nodes labelled 'visual.image' and 'text.embedding'
# Output nodes labelled 'results.prediction', 'results.features.[0]', 'results.features.[1]'
```

--------------------------------

### Create and Manage Tooltips

Source: https://github.com/sachinhosmani/torchvista/blob/main/torchvista/templates/graph.html

Provides functions to create, position, and remove custom tooltips for graph elements. Tooltips are dynamically generated and appended to the graph container, with options for different sizes.

```javascript
let currentTooltip = null; function createTooltip(text, smaller = false) { removeTooltip(); const tooltip = document.createElement('div'); tooltip.className = !smaller ? `custom-tooltip_${unique_id}` : `custom-tooltip_smaller_${unique_id}`; tooltip.textContent = text; const container = document.getElementById(`graph-container_${unique_id}`); container.appendChild(tooltip); currentTooltip = tooltip; return tooltip; }
```

```javascript
function positionTooltip(tooltip, buttonElement) { // Get the button's position in screen coordinates const buttonRect = buttonElement.getBoundingClientRect(); // Get the container's position in screen coordinates const containerRect = document.getElementById(`graph-container_${unique_id}`).getBoundingClientRect(); // Calculate position relative to container const buttonX = buttonRect.left + (buttonRect.width / 2) - containerRect.left; const buttonY = buttonRect.top - containerRect.top; // Show tooltip to measure its size tooltip.style.display = 'block'; const tooltipRect = tooltip.getBoundingClientRect(); // Position tooltip centered above button const left = buttonX - (tooltipRect.width / 2); const top = buttonY - tooltipRect.height - 6; // 6px gap above button tooltip.style.left = `${left}px`; tooltip.style.top = `${top}px`; }
```

```javascript
function showTooltip(buttonElement, text, smaller = false) { const tooltip = createTooltip(text, smaller); positionTooltip(tooltip, buttonElement); }
```

```javascript
function removeTooltip() { if (currentTooltip) { currentTooltip.remove(); currentTooltip = null; } }
```

--------------------------------

### Torchvista trace_model API Reference

Source: https://github.com/sachinhosmani/torchvista/blob/main/README.md

Overview of the parameters available for the trace_model function, including options for visualization, tracing depth, and export.

```python
trace_model(
    model,
    inputs,
    show_non_gradient_nodes=True,
    collapse_modules_after_depth=1,
    forced_module_tracing_depth=None,
    height=800,
    width=None,
    export_format=None,
    show_module_attr_names=False,
    export_path=None,
    show_compressed_view=False,
)
```

--------------------------------

### Initialize Collapsed Containers

Source: https://github.com/sachinhosmani/torchvista/blob/main/torchvista/templates/graph.html

Determines which containers should be initially collapsed based on their depth in the graph, ensuring repeat containers remain expanded.

```javascript
function initializeCollapsedContainersFromGraph() {
  const containerDepths = new Map();
  function getContainerDepth(container) {
    if (containerDepths.has(container)) {
      return containerDepths.get(container);
    }
    let depth = 0;
    for (const [child, parent] of Object.entries(ancestor_map)) {
      if (child === container && parent) {
        // Skip repeat containers when counting depth
        const parentDepth = getContainerDepth(parent);
        depth = isRepeatNode(parent) ? parentDepth : parentDepth + 1;
        break;
      }
    }
    containerDepths.set(container, depth);
    return depth;
  }
  for (const container of containerNodes) {
    const depth = getContainerDepth(container);
    if (depth >= collapse_modules_after_depth && !isRepeatNode(container)) {
      collapsedContainers.add(container);
    }
  }
  // Ensure repeat containers stay expanded
  for (const container of containerNodes) {
    if (isRepeatNode(container)) {
      collapsedContainers.delete(container);
    }
  }
}
```

--------------------------------

### trace_model

Source: https://github.com/sachinhosmani/torchvista/blob/main/docs/index.html

Traces a PyTorch model to generate an interactive visualization. It supports various options for controlling the visualization's appearance and export behavior.

```APIDOC
## trace_model

### Description
Traces a PyTorch model to generate an interactive visualization. It supports various options for controlling the visualization's appearance and export behavior.

### Method
`trace_model`

### Parameters
#### Parameters
- **model** (torch.nn.Module) - Required - Model instance to visualize.
- **inputs** (Any) - Required - Input(s) forwarded into the model; pass a single input or a tuple.
- **show_non_gradient_nodes** (bool, default: True) - Optional - Display nodes for constants and other values outside the gradient graph.
- **collapse_modules_after_depth** (int, default: 1) - Optional - Depth to initially expand nested modules; `0` collapses everything (nodes can still be expanded interactively).
- **forced_module_tracing_depth** (int, default: None) - Optional - Maximum depth of module internals to trace; `None` traces only user-defined modules.
- **height** (int, default: 800) - Optional - Canvas height in pixels.
- **width** (int | str, default: None) - Optional - Canvas width; accepts pixels or percentages; defaults to full available width when omitted.
- **export_format** (str, default: None) - Optional - Export format: `png`, `svg`, or `html` if exporting graph as a file. Otherwise, by default the graph is shown within the notebook.
- **show_module_attr_names** (bool, default: False) - Optional - Display attribute names for modules when available instead of just class names.
- **export_path** (str, default: None) - Optional - Custom path if exporting as a file. **Only HTML format** is currently supported with custom export paths.
- **show_compressed_view** (bool, default: False) - Optional - Compress the graph by showing repeating nodes of the same type with identical input and output dims in single "repeat" blocks. This feature currently only recognises repeating nodes within `Sequential` and `ModuleList`. Warning: this feature might be expensive on large models.
```

--------------------------------

### Label Nodes with Attribute Names using `show_module_attr_names`

Source: https://context7.com/sachinhosmani/torchvista/llms.txt

Set `show_module_attr_names=True` to label module nodes with their attribute names (e.g., `fc`, `embed`) instead of just their class names. This improves navigation in large models.

```python
import torch
import torch.nn as nn
from torchvista import trace_model

class Encoder(nn.Module):
    def __init__(self):
        super().__init__()
        self.embed    = nn.Embedding(100, 32)
        self.layer1   = nn.Linear(32, 64)
        self.layer2   = nn.Linear(64, 32)
        self.dropout  = nn.Dropout(0.1)

    def forward(self, x):
        x = self.embed(x)
        x = self.layer1(x)
        x = self.dropout(x)
        return self.layer2(x)

model  = Encoder()
inputs = torch.randint(0, 100, (2, 10))

# Default: nodes labelled "Embedding", "Linear", "Dropout", etc.
trace_model(model, inputs)

# With attr names: nodes labelled "embed", "layer1", "dropout", "layer2"
trace_model(model, inputs, show_module_attr_names=True)
```