v0.6.0: compiled forward/inverse maps + F/G sampler variants

Highlights

Compiled forward / inverse map fast paths (on ComposedTransform). Capture F = flow.t() and opt into torch.compile'd maps with one chainable call, mirroring Potential.enable_grad:

F = flow.t().enable_for_ladj().enable_inv_ladj()
y, ladj     = F.for_ladj(x)   # == F.call_and_ladj(x),     compiled
x_back, ilj = F.inv_ladj(y)   # == F.inv.call_and_ladj(y), compiled

Both return the fused (points, log|det J|). They live on ComposedTransform (zflows/core/transforms.py), not on Flow. They are deliberately not idempotent — re-calling enable_* recompiles against the current flow.t(), so in-place parameter updates (optimizer.step()) are tracked automatically while a .to(device)/dtype move just needs a fresh flow.t(). On an RTX 5070 Ti the forward sees ~3–20× and the spline-bisection inverse ~3–21× (largest at high d, where the raw inverse ~400 ms collapses to ~19–35 ms) — see the new benchmark tests/compare_compiled_inverse.py / writeup.

Flow-proposal annealed importance sampling. New annealed_importance_sampling_F / annealed_importance_sampling_G (with annealed_importance_sampling aliasing the _F forward-map variant) transport particles from exp(-source) to exp(-target) using a trained flow as the proposal, along the geometric path pi_k = mu_1^{k/M} (F_# mu_0)^{1-k/M}: per rung, reweight by the importance_weights_log rule, resample, then Langevin-rejuvenate in the target. The weights are computed directly from the raw potentials and the flow Jacobian (no bridge potential is built); the routine auto-uses the compiled for_ladj/inv_ladj maps and the target/source .eval fast paths when enabled. Algorithm and verification in tests/_verify_utils.py §F.

_F / _G variants for importance weights. importance_weights{,_F,_G} and importance_weights_log{,_F,_G} — the no-suffix name aliases the _F (forward map F) version; the _G twin takes the inverse map G = F^{-1} (target → source, e.g. flow.t().inv) and is mathematically identical. Both use the compiled for_ladj/inv_ladj maps when the passed transform has them enabled.

Breaking changes

• Loss functions renamed. source_KL_F → reverse_KL_F, target_KL_F → forward_KL_F, plus the inverse-map siblings reverse_KL_G / forward_KL_G. The reverse_KL / forward_KL aliases are unchanged (they still point to the _F versions), so existing training code is unaffected; only direct imports of the old source_KL_* / target_KL_* names must be updated.
• Potential.release() removed. Python refcounting + GC handle Potential lifetimes; the rarely-used manual teardown is gone.

Smaller additions

• Docker support: a Dockerfile and a Docker usage section in the README.
• README: documents the compiled for_ladj / inv_ladj flow feature, merges the auto-compiled gradient note (enable_grad / enable_eval) into the unified Potential section, adds a "one consistent interface, compiled or not" closing note, and adds the forward/inverse map benchmark (experiment 3) with its full results table.
• zflows/__init__.py public-surface list updated to the full {,_F,_G} families.

Verification

tests/_verify_potential.py, tests/_verify_flow.py (incl. §16 compiled-map checks), and tests/_verify_utils.py (incl. §D F/G + §F AIS) all pass on Linux + NVIDIA GPU.