Info

Last Execution: 2026-02-17

Package	Version
nnsight	0.5.15
Python	3.12.3
torch	2.10.0+cu128
transformers	5.2.0

Scan¶

model.scan() runs fake tensors through the model to collect shape and type information — without using any real memory or compute. This lets you inspect module dimensions and debug interventions before running them for real.

Setup¶

In [1]:

import nnsight
from nnsight import LanguageModel

model = LanguageModel("openai-community/gpt2", device_map="auto", dispatch=True)

import nnsight from nnsight import LanguageModel model = LanguageModel("openai-community/gpt2", device_map="auto", dispatch=True)

Getting Shape Information¶

Use model.scan() like model.trace(), but no real computation happens. You can inspect .shape on any module's output to learn its dimensions.

In [2]:

with model.scan("The Eiffel Tower is in the city of"):
    hidden_dim = nnsight.save(model.transformer.h[0].output[0].shape[-1])
    seq_len = nnsight.save(model.transformer.h[0].output[0].shape[1])
    vocab_size = nnsight.save(model.lm_head.output.shape[-1])

print(f"Hidden dim: {hidden_dim}")
print(f"Sequence length: {seq_len}")
print(f"Vocab size: {vocab_size}")

with model.scan("The Eiffel Tower is in the city of"): hidden_dim = nnsight.save(model.transformer.h[0].output[0].shape[-1]) seq_len = nnsight.save(model.transformer.h[0].output[0].shape[1]) vocab_size = nnsight.save(model.lm_head.output.shape[-1]) print(f"Hidden dim: {hidden_dim}") print(f"Sequence length: {seq_len}") print(f"Vocab size: {vocab_size}")

Hidden dim: 768
Sequence length: 10
Vocab size: 50257

You must save values to access them outside scan

model.scan() is a tracing context just like model.trace(). Values defined inside it are garbage collected when the block exits. Use nnsight.save() for non-tensor values like shape integers, or .save() for tensors.

Inspecting Shapes with Print¶

You can print() inside a scan context to inspect shapes interactively — useful for exploring an unfamiliar model.

In [3]:

with model.scan("The Eiffel Tower is in the city of"):
    print(f"Embedding output: {model.transformer.wte.output.shape}")
    print(f"Layer 0 output:   {model.transformer.h[0].output[0].shape}")
    print(f"Layer 11 output:  {model.transformer.h[11].output[0].shape}")
    print(f"LM head output:   {model.lm_head.output.shape}")

with model.scan("The Eiffel Tower is in the city of"): print(f"Embedding output: {model.transformer.wte.output.shape}") print(f"Layer 0 output: {model.transformer.h[0].output[0].shape}") print(f"Layer 11 output: {model.transformer.h[11].output[0].shape}") print(f"LM head output: {model.lm_head.output.shape}")

Embedding output: torch.Size([1, 10, 768])
Layer 0 output:   torch.Size([1, 10, 768])

Layer 11 output:  torch.Size([1, 10, 768])
LM head output:   torch.Size([1, 10, 50257])

Debugging Interventions¶

Scan catches shape errors before you run the real model. This is especially useful when working with large models where a failed run wastes significant time.

In [4]:

input_text = "The Eiffel Tower is in the city of"
num_tokens = len(model.tokenizer.encode(input_text))
print(f"Number of tokens: {num_tokens}")

# This will catch an off-by-one error
try:
    with model.scan(input_text):
        # Bug: num_tokens is 10, but valid indices are 0-9
        model.transformer.h[11].output[0][:, num_tokens, :] = 0
except IndexError as e:
    print(f"Caught error: {e}")

input_text = "The Eiffel Tower is in the city of" num_tokens = len(model.tokenizer.encode(input_text)) print(f"Number of tokens: {num_tokens}") # This will catch an off-by-one error try: with model.scan(input_text): # Bug: num_tokens is 10, but valid indices are 0-9 model.transformer.h[11].output[0][:, num_tokens, :] = 0 except IndexError as e: print(f"Caught error: {e}")

Number of tokens: 10
Caught error: 

Traceback (most recent call last):
  File "/tmp/ipykernel_253053/2942650569.py", line 9, in <module>
    model.transformer.h[11].output[0][:, num_tokens, :] = 0

IndexError: index 10 is out of bounds for dimension 1 with size 10

In [5]:

# Fixed version — index the last token correctly
with model.scan(input_text):
    model.transformer.h[11].output[0][:, num_tokens - 1, :] = 0
    print(f"Intervention shape check passed!")

# Fixed version — index the last token correctly with model.scan(input_text): model.transformer.h[11].output[0][:, num_tokens - 1, :] = 0 print(f"Intervention shape check passed!")

Intervention shape check passed!

Using Scan for Dynamic Dimensions¶

Scan is useful when you need to know a model's dimensions to construct tensors for interventions — like steering vectors or noise.

In [6]:

with model.scan("test"):
    dim = nnsight.save(model.transformer.h[0].output[0].shape[-1])

print(f"Hidden dimension: {dim}")

# Now use the dimension in a real trace
import torch

steering_vector = torch.randn(dim)

with model.trace("The Eiffel Tower is in the city of"):
    device = model.transformer.h[5].output[0].device
    model.transformer.h[5].output[0][:, -1, :] += steering_vector.to(device)
    logits = model.lm_head.output.save()

print(f"Steered prediction: {model.tokenizer.decode(logits[0, -1].argmax(dim=-1))}")

with model.scan("test"): dim = nnsight.save(model.transformer.h[0].output[0].shape[-1]) print(f"Hidden dimension: {dim}") # Now use the dimension in a real trace import torch steering_vector = torch.randn(dim) with model.trace("The Eiffel Tower is in the city of"): device = model.transformer.h[5].output[0].device model.transformer.h[5].output[0][:, -1, :] += steering_vector.to(device) logits = model.lm_head.output.save() print(f"Steered prediction: {model.tokenizer.decode(logits[0, -1].argmax(dim=-1))}")

Hidden dimension: 768

Steered prediction:  London

When to use scan

• Exploring unfamiliar models — quickly check shapes at every layer without running the full model
• Debugging interventions — catch shape mismatches and index errors before a costly real run
• Dynamic tensor construction — learn hidden dimensions to build correctly-sized steering vectors, probes, or adapters
• Remote execution — validate interventions locally before sending them to NDIF