tdhook 🤖🪝

Composable interpretability for PyTorch with TensorDict and torch hooks.

Getting Started

Install TDHook from PyPI:

pip install tdhook

TDHook methods wrap an ordinary PyTorch model for the lifetime of a context manager. Inputs, baselines, model outputs, and interpretability results use explicit TensorDict keys:

import torch
from torch import nn
from tensordict import TensorDict
from tdhook.attribution import IntegratedGradients

model = nn.Sequential(nn.Linear(4, 8), nn.ReLU(), nn.Linear(8, 2))
inputs = torch.tensor([[0.2, -0.1, 0.4, 0.7]])

def select_score(outputs, _):
    score = outputs["output"][..., 0]
    return TensorDict(score=score, batch_size=outputs.batch_size)

data = TensorDict(
    {
        "input": inputs,
        ("baseline", "input"): torch.zeros_like(inputs),
    },
    batch_size=[1],
)

with IntegratedGradients(init_attr_targets=select_score).prepare(model) as hooked_model:
    result = hooked_model(data)

attributions = result["attr", "input"]

The context installs and removes the hooks; the returned attribution has the same shape as inputs. See Getting Started for the annotated version.

Learn by example

The tutorial gallery collects all maintained method and end-to-end notebooks. Launch a method notebook directly in Colab:

• Integrated Gradients:
• Steering Vectors:
• Linear Probing:
• Bilinear Probing:
• Dimension Estimation:
• Representation Similarity:

Use the generated API reference for exact signatures. The TDHook agent skill provides guidance for attribution, activation analysis, probing, steering, and weight-level interventions.

Config

This project uses uv to manage python dependencies and run scripts, as well as just to run commands.

Benchmarks

The maintained benchmark suite checks current TDHook attribution, capture, and intervention behavior against reference libraries before recording versioned timing and memory results. It provides a cheap local smoke mode and a documented full mode; it does not claim to reproduce the historical v0.1 paper measurements.

Citation

If you're using tdhook in your research, please cite it using the following BibTeX entry:

@misc{poupart2025tdhooklightweightframeworkinterpretability,
      title={TDHook: A Lightweight Framework for Interpretability},
      author={Yoann Poupart},
      year={2025},
      eprint={2509.25475},
      archivePrefix={arXiv},
      primaryClass={cs.AI},
      url={https://arxiv.org/abs/2509.25475},
}

License

tdhook is licensed under the MIT License. See LICENSE for details.