mixle.inference.jit module

JIT-compile a model’s whole-tree log-density to a single XLA program (A2).

mixle’s per-distribution backend_seq_log_density is engine-neutral: a Composite/Mixture/… combines its children’s scores through ComputeEngine ops, so the entire model tree is one pure array computation. Run it on the JAX engine under jax.jit and the whole tree lowers to a single compiled XLA program – the “compile the model once” trick that makes JAX-backed PPLs fast.

score = jit_seq_log_density(model) # compiles lazily, per data shape ll = score(data) # bit-identical to model.seq_log_density(encode(data)), faster

This is ideal for repeated scoring of a fixed model – prediction, held-out / cross-validation log-likelihoods, importance weights, bootstrap, large-batch scoring – where the same compiled program is reused across calls. The result is bit-identical to model.seq_log_density (verified) and, on a large composite tree, several-fold faster than vectorized NumPy.

Scope note: this compiles with the model’s parameters baked into the program (the structure lowers to one XLA function). Reusing a single compiled program ACROSS EM iterations as the parameters update needs the parameters threaded as traced inputs (a tree-level backend_log_density_from_params); that is the next increment and is not done here.

class JittedScorer(model, engine=None)[source]

Bases: object

A jax.jit-compiled whole-tree log-density scorer for a fixed model.

Calling the scorer encodes data, runs the engine-neutral backend_seq_log_density over the whole model tree on the JAX engine under jax.jit, and returns host log-densities. The compiled program is cached and reused across calls with the same data shape.

Parameters:

  • model (Any)

  • engine (Any)

jit_seq_log_density(model, engine=None)[source]

Return a JittedScorer: the whole-tree log-density of model compiled to one XLA program.

jit_seq_log_density(model)(data) is bit-identical to model.seq_log_density(encode(data)) but runs as a single jax.jit XLA program over the entire composite tree – fast for repeated scoring of a fixed model. Requires the JAX optional extra and that every leaf in model declares JAX support (raises EngineNotSupportedError-style errors from the engine layer otherwise).

Parameters:
Return type:

JittedScorer

jit_em_mixture(model, data, *, max_its=100, engine=None)[source]

Fit a finite mixture of same-family scalar exponential-family leaves by EM, with the ENTIRE EM loop (every E-step + closed-form weighted M-step iteration) compiled to ONE jax.jit XLA program via lax.scan – the parameters are traced inputs threaded through the loop on-device, so there is no per-iteration recompile and no per-iteration host sync. This is roadmap A2 bullet 1 (“repeated EM iterations run as one XLA program”) realized literally.

model is the initial mixture (its components seed the EM); supported leaves: Gaussian, Poisson, Exponential. Runs a fixed max_its iterations (no host-side early stop – that is the point: the loop stays on-device). Returns a fitted MixtureDistribution, bit-close to the host EM from the same start (it is the same EM update). Raises NotImplementedError for unsupported structure.

SPEED – measured, honest: the payoff is GPU/TPU and large scale, where XLA parallelizes the E-step over millions of points and many components. On an Apple M4 GPU (via jax-metal) this kernel runs ~21x faster than mixle’s vectorized NumPy EM (K=10, N=1e6, 50 iters: 156 ms vs 3254 ms) with identical estimates, and ~12x faster than the same jitted loop on CPU. On CPU it is *not* a speedup – mixle’s host EM is already vectorized NumPy (+ a fused path) and a long sequential scan of small steps loses to it (~0.1-0.8x for K up to 40). So: GPU -> big win, CPU -> use the host EM. (The CPU win from A2 is the single-pass scoring jit, jit_seq_log_density(), ~8x.) Note: jax-metal is version-pinned – the GPU result above used Python 3.11 + jax/jaxlib 0.4.34 + jax-metal 0.1.1; newer jaxlib emits StableHLO that jax-metal 0.1.1 cannot compile.

Parameters: