v0.5.0

Highlights

Three independent extensions to the v0.4.0 API surface, all backward-compatible at their default values.

Inverse-temperature beta everywhere it has a unique physical meaning. A single beta knob (defaulting to 1.0) joins the four KL losses, langevin / hmc, and Gaussian.samples; importance_weights{,_log} get two betas (beta_source / beta_target) for tempered SMC ladder rungs.

loss.compile split into compile_raw / compile_beta.

• compile_raw(loss_fn, *captured) — single-input fast path, returned closure is (x_batch) -> scalar, every non-batch argument baked in as a Python closure constant. No wrapper / no per-call as_tensor, lowest-overhead choice for fixed-beta training.
• compile_beta(loss_fn, *captured) — same captured-once trick but the closure exposes (x_batch, beta=1.0) -> scalar. A 0-d-tensor cast at the boundary keeps Dynamo from value-specializing on the Python float, so a single compiled artifact handles every value of beta. Use for annealed / adaptive schedules.
• The old loss.compile name is gone; both helpers live in zflows.loss and can be imported directly: from zflows.loss import reverse_KL, compile_raw.

RealNVP learnable linear mixing (Glow-style 1×1 conv). New optional mixing: str | None = None kwarg interleaves a learnable $d \times d$ map between every two consecutive couplings. mixing=\"rotation\" is an orthogonal $R = \exp(A - A^T)$ (log-det $\equiv 0$); mixing=\"lu\" is a PLU map whose learnable $L$-diagonal supplies a non-trivial log-det. Both initialise at identity. The default None keeps the byte-exact pre-mixing architecture.

Smaller additions

• Linear_Combination generalised from the two-potential form to an N-potential form: Linear_Combination([U_0, U_1, ...], [c_0, c_1, ...]). Coefficients can also be a 1-d torch.Tensor, registered as an nn.Module buffer so .to(device) follows them in lockstep with the potentials.
• New utility helpers in zflows.utils: suppress_warnings(), set_cache_size_limit(N), check_compile_available() — three one-shot calls to silence Triton / Inductor / Dynamo / Python chatter, give Dynamo more cache headroom, and verify the compile toolchain.
• randmask: bool = True is now the default per-layer feature ordering for NSF / NCSF / RealNVP (fresh torch.randperm(d) each layer); randmask=False switches back to the legacy arange / arange.flip alternation.
• Gaussian.samples(N, beta=1.0) draws from the tempered $\mathcal N(\mu, \Sigma / \beta)$.

Documentation

• README gains a Guidance on choosing flow classes subsection (comparison table of NSF / NCSF / RealNVP / CNF + decision-tree code snippet), a section on the unified Potential class with the super().__init__() safety warning, install step 4 for check_compile_available(), and FAQ entries on conditional flows and JAX alternatives.
• New benchmark writeup tests/compare_compiled_loss.md: compile_raw vs. raw reverse_KL across a dimension × hidden_features grid on an RTX 5070 Ti — 4–10× per-step speedup with mode='default', up to ~13× with mode='reduce-overhead'.

Verification

The test suite was consolidated into three banner-separated harnesses, all CPU-runnable, no external fixtures:

• tests/_verify_flows.py — 14 sections covering construction, round-trip bijection, .zeros() identity, call_and_ladj vs autograd slogdet, NSF/NCSF box preservation, RealNVP forward+inverse ladj cancellation, CNF Hutchinson unbiasedness, the new RealNVP mixing layers, and the compile_raw / compile_beta cache-stability check.
• tests/_verify_potential.py — Potential.enable_eval / Linear_Combination / Gaussian.samples(beta) checks.
• tests/_verify_utils.py — tamed Langevin, HMC, L-BFGS, and the new beta-tempered-stationary / IS-linearity-in-beta checks.

Plus a tests/_all.py runner that executes every demo under MPLBACKEND=Agg.

Dependencies

Unchanged: `torch>=2.10.0`, `numpy>=2.4.0`. `numpy` is only used in `zflows/core/numerics.py` for a one-shot `np.polynomial.legendre.leggauss` call (cached), so the runtime footprint is the same as v0.4.0.

🤖 Generated with Claude Code