FlatPPL for Python¶
Load or write a FlatPPL model, compile its density, and pass it to BlackJAX. FlatPPL supplies the model and compiled functions. BlackJAX supplies inference.
The package bundles the Rust compiler and its StableHLO emitter. JAX/XLA executes the emitted program. Enzyme-JAX supplies differentiation.
Install with Pixi, compile your first function, or fit a mixed model with BlackJAX.
Sample in unconstrained coordinates¶
For a scalar log-density query, compile the coordinate transforms with the target:
target = query.compile_transformed(logdensity=True)
initial = target.unconstrain(point)
# Pass target.unconstrained and initial to BlackJAX.
# Convert retained draws back with target.constrain.
Rust chooses the transforms from the query’s input domains. Positive scales, bounded parameters, ordered cutpoints, simplexes and zero-sum effects use their supported canonical charts. The density correction runs inside StableHLO. You do not write Jacobians or call Python during sampling.
The mixed-model tutorial shows the complete four-chain workflow, diagnostics and prediction. The transform reference lists the supported domains and limits.
Models and queries use the same Module type. You choose their names and roles.
Use query.compile() when you want the declared, constrained input interface.
The alpha package supports Python 3.12–3.15 on Linux x86-64 and Apple silicon macOS. See installation for the exact platform requirements and JAX support for differentiation limits.
Get started