Composable interpretability for PyTorch with TensorDict and torch hooks.
Install TDHook from PyPI with Python 3.11 or newer:
pip install tdhookTDHook 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.
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.
This project uses uv to manage python dependencies and run scripts, as well as just to run commands.
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.
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},
}
tdhook is licensed under the MIT License. See LICENSE for details.