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, and 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

The package is on PyPI as tramdag; uv add tramdag or pip install tramdag installs it, and the plots extra (tramdag[plots]) adds matplotlib for the figures. The PyPI release lags main; for the current state install from git (tramdag @ git+https://github.com/tensorchiefs/tramdag.git@main, pinned to a commit for reproducibility). A clone with uv sync gives the development setup with tests and experiments.

Documentation

Guides, executed notebooks and the API reference are published at the documentation site. Each subject has one guide, and its worked example is one notebook.

CONTRIBUTING.md covers setup, tests, linting and releases; notebooks/README.md how the notebooks are written and run.

Experiments

experiments/ replicates the paper's experiments against its own R code and checks them against committed ground truth in continuous integration. experiments/README.md explains the layout, the YAML variants and the frozen data; docs/paper-replication.md gives the protocol, the hyperparameters and the numbers, experiment by experiment.

Layout

• src/tramdag/ is the framework, and nothing else.
• tests/ measures it against three inline data-generating processes; it never imports research code.
• experiments/ holds the paper replications and their frozen data sets.
• notebooks/ are jupytext .py files, executed by the docs workflow.
• docs/ are the guides; the API pages render from the docstrings.

Use of AI tools

This project was developed with the help of LLM coding assistants. The authors reviewed and take responsibility for all content.

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},
}