Phantom types for robust hierarchical models with typegeist

TL;DR

This paper introduces typegeist, a Python type system that uses static analysis to enforce correct data-parameter indexing in Bayesian hierarchical models, reducing bugs and ensuring model correctness.

Contribution

It presents a novel type system for Python that improves robustness of hierarchical models by statically verifying data-parameter correspondences.

Findings

Typegeist effectively detects indexing errors in probabilistic models.

The system integrates with existing probabilistic programming frameworks.

Minimal runtime overhead is introduced by the static analysis.

Abstract

In Bayesian hierarchical models, group-level parameter arrays must be mapped to the observation axis, often using explicit indexing. In complex models with numerous incompatible data and parameter sets, this introduces the potential for bugs, as indexing with the incorrect indices typically fails silently. Here we present typegeist, a type system for Python that uses static type analysis to enable specification and enforcement of data-parameter-index correspondences. We show how this can be used with common probabilistic programming frameworks to help guarantee model correctness with minimal run-time overhead.