tramdag: Interpretable Neural Causal Models (TRAM-DAGs) in PyTorch

TRAM-DAGs model each variable of a structural causal model with a monotone normalizing flow, the statistician's transformation model. The whole model is one triangular normalizing flow from iid standard-logistic latents to the observed variables. Its triangular structure is exactly your causal DAG. One fit on observational data answers all three rungs of Pearl's causal hierarchy: observational, interventional (the do-operator) and counterfactual (abduction). The model keeps interpretable effects: every linear-shift coefficient is a log-odds ratio, as in a classical proportional-odds model.

Beate Sick & Oliver Dürr, Interpretable Neural Causal Models with TRAM-DAGs, CLeaR 2025 (arXiv:2503.16206). This repository is the reference implementation, in PyTorch on zuko. It replicates the paper's experiments under continuous integration.

Start here

The demo notebook is the introduction. It fits the paper's bimodal benchmark, walks the three rungs, and checks every answer against analytic ground truth. It runs on CPU in a few minutes, in Colab or locally.

Install

uv add tramdag or pip install tramdag installs the package from PyPI. The plots extra (tramdag[plots]) adds matplotlib for the figures. A clone with uv sync gives the development setup (CONTRIBUTING.md).

Documentation

The documentation site publishes the guides, the executed notebooks and the API reference.

• The model: docs/model.md and docs/notation.md
• Reading a fitted model: docs/interpretation.md
• Fitting: docs/fitting.md
• Treatment effects that vary with covariates: docs/varying-coefficients.md and docs/scores.md
• The design of the package: docs/architecture.md and docs/code-map.md
• The paper replications: docs/paper-replication.md and experiments/README.md
• The tests: tests/README.md
• The notebooks: notebooks/README.md

Use of AI tools

LLM coding assistants helped to develop this project. The authors reviewed all content and take responsibility for it.

Citation

If you use tramdag, please cite the method paper:

@inproceedings{sick2025tramdag,
  title     = {Interpretable Neural Causal Models with {TRAM-DAGs}},
  author    = {Sick, Beate and D{\"u}rr, Oliver},
  booktitle = {Proceedings of the 4th Conference on Causal Learning and Reasoning (CLeaR)},
  series    = {Proceedings of Machine Learning Research},
  volume    = {275},
  pages     = {606-630},
  year      = {2025},
  eprint    = {2503.16206},
  url       = {https://arxiv.org/abs/2503.16206},
}