trace
Run a forward pass on an open-weight language model to identify the attention heads and MLP neurons most responsible for its predicted next token, explaining why the model made that prediction.
Instructions
Run a forward pass of a small-to-medium open-weight language model (via TransformerLens) on one prompt, and report which attention heads and MLP neurons were most responsible for its predicted next token: heads ranked by direct logit attribution, neurons ranked by activation magnitude at the final prompt position. Call this to answer 'why did the model predict X' for a specific prompt. It only works on models TransformerLens's HookedTransformer.from_pretrained supports (GPT-2, Pythia, Llama, Gemma, Qwen, and similar open-weight checkpoints), not closed-source APIs like OpenAI or Anthropic models. Read-only and deterministic for a given model, prompt, and top_k: it writes nothing except the model's own weights, which HuggingFace Hub downloads to a local cache (~/.cache/huggingface) the first time a given model name is requested (needs network access that one time; later calls for the same model run offline from cache). Runs on CPU by default and can be slow for large models. On failure (an unsupported model name, or a prompt longer than the model's context window) it returns a structured error object instead of raising, so the tool call itself never fails silently. Parameters: model (str) is any name HookedTransformer.from_pretrained accepts, e.g. 'gpt2' or 'EleutherAI/pythia-70m'; prompt (str) is the input text; top_k (int, default 10) caps how many top heads and neurons are returned. Example call: model='gpt2', prompt='The capital of France is Paris. The capital of Japan is', top_k=5. Returns JSON with schema_version, operation, model (resolved name, backend, device, and layer/head/dimension counts), prompt, predicted_token, predicted_token_id, top_neurons (list of {layer, neuron_index, activation}), and top_heads (list of {layer, head_index, logit_attribution}).
Input Schema
| Name | Required | Description | Default |
|---|---|---|---|
| model | Yes | ||
| top_k | No | ||
| prompt | Yes |
Output Schema
| Name | Required | Description | Default |
|---|---|---|---|
No arguments | |||