Releases: xuda-ye-math/zflows
Release list
v0.6.2
v0.6.1: sequential_monte_carlo + hamiltonian_monte_carlo + nested linear_combination
A backward-compatible follow-up to v0.6.0 (no breaking changes — hmc keeps working via an alias).
New
sequential_monte_carlo — flow-free annealed-Langevin SMC. Transports particles from exp(-beta_source * source) to exp(-beta_target * target) (same space, no flow) through a ladder of linearly-interpolated bridge potentials, alternating reweight + multinomial resample + Langevin on each bridge u_k:
samples, ess = sequential_monte_carlo(x0, source, target,
beta_source=1.0, beta_target=1.0,
ladder=20, step=5e-3, iters=30, chunk=1)The bridge u_k = (1 - k/M)*beta_source*source + (k/M)*beta_target*target is built once with linear_combination and retuned per rung via set_coeffs. It returns (samples, ess), where ess is a list[float] of the per-rung effective sample size (pre-resample) — a ladder-spacing diagnostic. This differs from annealed_importance_sampling_F/_G, which use a trained flow as the proposal and rejuvenate only in the target. The per-rung reweighting uses the compiled .eval fast path when enable_eval() is set.
Changed
hmc → hamiltonian_monte_carlo (clearer canonical name); hmc is kept as an alias, so existing code is unaffected. The optimization = lbfgs / rejuvenation = langevin aliases now sit directly after their function definitions.
Nested linear_combination is supported. The previous "do NOT nest" caveat is removed: a linear_combination may be a child of another — the combined .grad / .eval closures compose recursively, enable_grad() / enable_eval() cascade to all leaves, and the fresh-allocation buffer-aliasing safety holds at every level (correct under both mode='default' and mode='reduce-overhead').
Tests & docs
- New
tests/_verify_nested.pyexercises nested combinations end-to-end (depths 2–3, enable cascade,set_coeffspropagation, tensor coeffs, default==reduce-overhead, runtime gating, and MALA/HMC reproducing an analytic Gaussian) — focused on the compiled grad/eval fast paths. - README: a warning that the compiled fast paths are not a free lunch — at high
dwith large conditioners (e.g.d=256,hidden_features=(256,256)) compiling theNSFfor_ladj/inv_ladjmaps is extremely slow to compile, with no effective fix (dropenable_*and use the identical plain maps) — plus a "Why istorch.compileso slow?" FAQ entry.
Verification
tests/_verify_potential.py, _verify_flow.py, _verify_utils.py, and _verify_nested.py all pass on Linux + NVIDIA GPU.
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), compiledBoth 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 siblingsreverse_KL_G/forward_KL_G. Thereverse_KL/forward_KLaliases are unchanged (they still point to the_Fversions), so existing training code is unaffected; only direct imports of the oldsource_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
Dockerfileand a Docker usage section in the README. - README: documents the compiled
for_ladj/inv_ladjflow feature, merges the auto-compiled gradient note (enable_grad/enable_eval) into the unifiedPotentialsection, 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__.pypublic-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.
v0.5.6: lowercase factories + auto-compiled linear_combination
Highlights
Three coupled API moves that simplify how user code talks to zflows.potential, all naturally backward-incompatible at the import line.
linear_combination lowercase factory. New linear_combination(potentials, coeffs=None) -> Linear_Combination returns a ready-to-use instance — the project-wide convention is uppercase = class, lowercase = instance, and call-sites that use the combination should look like instance-creation, not class-construction. The Linear_Combination class is still exported for isinstance(u, Linear_Combination) checks and subclassing. Every user-facing example in README.md, every test script, and every writeup under tests/*.md now uses the lowercase form.
potential_from(fn) returns an instance. The old potential_from returned a Potential subclass that the caller then instantiated; the companion potential_instance_from was a one-line wrapper around potential_from(fn)(). Both collapse into a single factory: potential_from(fn) now returns the instance directly, so the canonical line is u = potential_from(myforward).to(device). potential_instance_from is removed; callers using it must rename to potential_from.
Linear_Combination auto-inherits its children's compiled .grad / .eval fast paths. The combined _grad_fn / _eval_fn are linked at __init__ time as Python closures sum_k self.coeffs[k] * U_k.grad(x) (and same for .eval(x)). They read self.coeffs[k] fresh on every call, so:
- Enabling the children once is enough (
u0.enable_grad(),u1.enable_grad()) —linear_combination([u0, u1], [c0, c1]).grad(x)then works immediately, with nolc.enable_grad()call needed. -
set_coeffs([c0_new, c1_new])is a pure coefficient update: no compiled artifact is invalidated, no Dynamo recompile is triggered. Annealed-bridge workflows (e.g.tests/4D_Boltzmann_generator.py) now pay each child'svmap(grad(forward))andcompile(forward)cost exactly once for the entire$M$ -rung anneal. - The "needs
.enable_grad()first" gate moves from a pre-flight check onLinear_Combinationto a runtime check at the child level (an un-enabled child raisesRuntimeErrorthe first time the closure tries to call its.grad(x)).
Smaller additions
- Device-consistency check at
Linear_Combinationconstruction. Scans everyparameter()andbuffer()across the child potentials; raisesRuntimeErrorif more than one device shows up, with the offending list spelled out. Pure-function children (e.g.potential_from(fn)instances,Uniform) contribute no device and are silently allowed. set_coeffs(coeffs)simplified. The previousenable_grad: bool = False/enable_eval: bool = Falseflags are gone — under the new closure-based design they were no-ops on hot instances and confusing on cold ones. Pure coefficient update only.Linear_Combination(potentials, coeffs=None)defaults to uniform 1/N ifcoeffsis omitted.coeffsis always stored as a plainlist[float]; tensor inputs are detached and.tolist()-converted at construction, so there's no buffer-registration /requires_gradsubtlety.- Docstring caveats on
loss_compile/loss_compile_beta. These helpers have only been tested againstNSF/NCSF. Behaviour onCNF,OTFlow,RealNVP, or any other flow class is unknown — the docstrings now flag this and recommend verifying against the un-compiledreverse_KL(x, target, flow.t())baseline before relying on the speedup. - Docstring caveat on nesting
Linear_Combination. Using aLinear_Combinationinstance as a child of another is untested and not recommended; the buffer-aliasing assumption (every child's compiled.grad(x)returning a static reduce-overhead buffer that's consumed before the next child runs) has not been exercised under nesting. List every constituent potential in one flatlinear_combination([u_0, ..., u_N], [c_0, ..., c_N])call instead.
Documentation
- README's "Auto-compiled gradient feature" section reordered: the leaf-
Potential.enable_grad()story is introduced first, thenlinear_combinationis shown inheriting it. The duplicate "Precompiled gradients on Potential" sub-header is gone. - README user-facing examples scrubbed for the
U= class /u= instance convention: the user-subclass example usesclass MyPotential(Potential)+u = MyPotential().to(device); the factory example usesu = potential_from(myforward).to(device)(the function ismyforward, notU_fn); thelinear_combinationexample usesu0,u1,u. - ESS narrative for the 4D Boltzmann generator refreshed against this release's run:
$\sim 0.68$ at$k=1$ , jumps straight to the$0.96$ –$0.97$ band by$k=2$ , holds there for the remaining ten rungs. - Compiled-loss benchmark table refreshed from the new
tests/compare_compiled_loss.csv: same RTX 5070 Ti hardware, $\sim$5–13 ms raw / 1.3–1.7 msmode='default'/ 0.5–1.7 msmode='reduce-overhead', 4–9× / 7–12× speedups respectively. - Test writeups (4D BG, 3D periodic, 2D reverse-KL, …) updated to describe the new factory pattern and the refreshed line-number pointers.
Tests
- Stateless
Potentialsubclasses →potential_from(forward_fn).tests/2D_reverse_KL.py'sclass U1,tests/_verify_potential.py'sclass U0/class U1,tests/_verify_utils.py'sclass Quartic/class Quartic1Dall converted. Stateful subclasses with constructor args —U_targetintests/4D_Boltzmann_generator.py(r0, a, q2, eps),U_targetintests/3D_periodic.py(kappa, renamed fromU1),Multi_Wellintests/multi_well_compare.py— kept as classes sincepotential_fromonly wraps stateless callables. u0/u1→u_source/u_targetacross every relevant test (2D forward-KL, 2D reverse-KL, 2D two-moons, 2D RealNVP, 3D periodic,_verify_potential). The bridge-instance naming now matches what 4D BG has been using since v0.5.5.Linear_Combination(...)→linear_combination(...)at every call-site in the test suite.- New
tests/_verify_potential.pysection B.5 covers the new closure-based behaviour: combined_grad_fn/_eval_fnlinked at__init__, runtime gate at the child level,set_coeffspreserves everyid()and the closure picks up the new coefficients with err = 0.0 against autograd / forward,enable_grad/enable_evalpropagate unconditionally, hot children survive via child-level idempotency, and alinear_combinationof pre-enabled children works immediately withoutlc.enable_grad().
Verification
Full sweep via python -m tests._all: 11 passed, 0 failed in 677.5 s on the same RTX 5070 Ti. The committed tests/_all.out records the per-script wall time (62.8 s _verify_flow, 16.1 s _verify_potential, 7.2 s _verify_utils, 254.8 s compare_compiled_loss, 257.2 s multi_well_compare, etc.) and full stdout.
Migration
```python
v0.5.5
from zflows.potential import Linear_Combination, potential_from, potential_instance_from
U = potential_from(my_fn) # returned a class
u = U().to(device) # had to instantiate yourself
u_inst = potential_instance_from(my_fn).to(device) # one-liner alternative
lc = Linear_Combination([u0, u1], [c0, c1])
lc.enable_grad() # had to enable on the combination
v0.5.6
from zflows.potential import linear_combination, potential_from
u = potential_from(my_fn).to(device) # returns instance directly
u0.enable_grad(); u1.enable_grad() # enable on the children
lc = linear_combination([u0, u1], [c0, c1]) # lc.grad(x) works immediately
```
Dependencies
Unchanged: torch>=2.7, numpy>=2.0. No new runtime imports.
🤖 Generated with Claude Code
v0.5.4
Packaging-only release. Fixes the PyPI page so figures actually render, and updates the install instructions to lead with the now-published `pip install zflows`.
README
- Rewrote 7 `<img src="...">` paths from repo-relative (`logo.png`, `tests/*.png`) to absolute `https://raw.githubusercontent.com/xuda-ye-math/zflows/main/...\` URLs. PyPI's Markdown renderer doesn't fetch repo-relative references, so the v0.5.3 PyPI page shipped with broken-image placeholders. Absolute `raw.githubusercontent` URLs work on both GitHub and PyPI — same pattern zuko uses.
- Installation: leads with `pip install zflows` for the released PyPI path; the existing `git clone` + editable-install workflow now follows as the "latest unreleased features + demo scripts on disk" alternative.
No library code changes
`state_dict` keys and runtime behaviour byte-identical to v0.5.3.
🤖 Generated with Claude Code
v0.5.3
Packaging-only release: prepares the PyPI page so the first upload through the new tag-driven workflow lands with a complete project description, classifiers, and URL block.
pyproject.toml polish
readme = \"README.md\"— PyPI renders the README as the long description.- Author block now carries the maintainer email.
keywords: `normalizing-flows`, `boltzmann-generator`, `pytorch`, `sampling`, `smc`, `neural-spline-flows`, `ot-flow`.- 12 trove
classifiers: Development Status, Intended Audience, License, Linux, Python 3.10–3.14, Scientific/Engineering :: AI / Mathematics / Physics. [project.urls]Homepage / Repository / Issues / Changelog all point at the GitHub repo + releases page.
CI / publish workflow
Already on main from b31bd31: `.github/workflows/publish-to-pypi.yml` triggers on every `v*` tag push, builds sdist + wheel with `python -m build`, lints via `twine check`, and uploads to PyPI through OIDC trusted publishing — no API token in repo secrets, gated by the `pypi` environment for manual approval.
README
Stale "Can I install zflows from PyPI?" FAQ entry removed.
No code changes
Library, tests, and demos are byte-identical to v0.5.2. This release exists solely to wire the metadata + workflow so PyPI can serve a complete page on the first upload.
🤖 Generated with Claude Code
v0.5.2
Patch release rolling up everything since v0.5.1 — one new flow class, two API additions, two correctness fixes, and a README cleanup pass.
New: OTFlow (optimal-transport continuous flow)
zflows.flow.OTFlow ports OT-Flow (Onken–Fung–Li–Ruthotto, AAAI 2021) into the same lazy-flow.t() contract as the other flows. The ODE velocity field is parameterised as the negative gradient of a learnable scalar potential Φ_θ (`OTPhi` — antiderivative-of-tanh ResNet + low-rank quadratic head), which makes the divergence tr(∇²Φ) closed-form O(d·m) — no Hutchinson estimate, no augmented Jacobian ODE. Integration is fixed-step RK4 (same RK4 used by CNF now), so per-step cost is deterministic and torch.compile-friendly.
OTFlow plugs into reverse_KL, compile_raw, compile_beta, the IS / Langevin / HMC utilities, and the capture-once contract — tests/_verify_flows.py §10/§14g/§15 confirm autograd-slogdet agreement, OT regularizer integration, and Dynamo cache stability under capture-once.
New: Potential._from(fn) factory
Wrap a stateless callable as a Potential in one line — no subclass boilerplate:
```python
def U(x):
return 0.5 * (x ** 2).sum(-1) + 2 * torch.cos(x[:, 0])
u = Potential._from(U).to(device).enable_grad()
g = u.grad(x)
```
The wrapped Potential supports the full toolchain (.to, .enable_grad, .enable_eval, .parameters). Method name is _from because from is a Python keyword. Stateful potentials (with constructor args, learnable sub-modules, etc.) still want to subclass Potential directly.
API improvements
check_compile_available()now searches/usr/local/cuda/bin/nvccand/usr/local/cuda-*/bin/nvcc(Ubuntu's default CUDA-Toolkit install locations, often not on $PATH) as fallbacks whenshutil.which(\"nvcc\")misses. Versioned directories are tried newest-first. The warning text enumerates every location searched.GeneralCouplingTransformgained a.zeros()method that does the last-conditioner-MLP weight/bias zero-out previously inlined inRealNVP.zeros(). The latter collapsed to `for layer in self._layers: layer.zeros()` — same iteration handles couplings and mixing layers uniformly.
Correctness fixes
LULinearTransformnumerical safety. `log|det| = sum log|diag(L)|` was uncovered: nothing in the parameterisation kept the LU diagonal away from zero, so a long training run could drive an entry across zero and explode the loss to NaN. `log_abs_det_jacobian` and `call_and_ladj` now clamp|diag(L)|to a1e-12floor before taking the log — ladj stays finite, gradient saturates to 0 on the affected entry. HealthyLU = Istill givesladj = 0byte-exact. Docstring updated; the principled Glow-stylesign(s0) · exp(log_s)re-parameterisation is noted as a future-major-release option (it would changestate_dictlayout).compile_betascalar-only assertion. The wrapper nowassert beta.dim() == 0to reject shape-[1]/[N]tensors. Higher-rank betas would silently trigger per-shape recompiles, breaking the central "one compiled artifact handles every beta" contract. Docstring tightened to "scalarfloat | 0-d Tensor" with an explicit note that per-sample tempering needs a different helper.Linear_Combinationgrad-tracking-coeff guard. A user passing `torch.tensor([...], requires_grad=True)` ascoeffswould have it registered as a buffer — silently hidden fromoptimizer.parameters()while backward still accumulated.grad. The constructor now `assert not coeffs.requires_grad` with a clear pointer tocoeffs.detach()or the explicit-nn.Parameterescape hatch.
Tests
tests/_verify_flows.py and tests/_verify_potential.py gained four new sub-blocks:
- §14h —
LULinearTransformnear-singular safety (diag(L) = 1e-20 still gives finite ladj). - §15 — captured-
Fparameter-mutation consistency on NSF + CNF + OTFlow. - §14g —
OT_loss = reverse_KL + OT regularizerschannel agreement. - §B.3 —
Linear_Combinationrejectsrequires_grad=Truetensor coeffs.
All three verify suites (_verify_flows.py, _verify_potential.py, _verify_utils.py) pass clean.
Documentation
- Features section: condensed the continuous-flows paragraph (dropped
$O(d \cdot m)$ log-det math + augmented-Jacobian-ODE comparison), the "swapping flow classes" paragraph, and thesuper().__init__()warning. Net result: ~3x shorter, same information for first-time readers. - Added a
Potential._from(fn)code block to the Potential section. - Added a
zflows/package-layout tree at the end of the Features section. - FAQ:
check_compile_availablebullet updated to mention the new fallback search paths. - Install step 4 (optional
check_compile_available()smoke test) added.
Compatibility
`state_dict` keys unchanged across all flow / potential classes. Code following the documented flow.t() → ComposedTransform capture-once pattern continues to work exactly as before. Existing v0.5.1 saved models load into v0.5.2 with no migration.
🤖 Generated with Claude Code
v0.5.1
Bug fix
RotationTransform / LULinearTransform silently broke the capture-once contract. Both classes ported from zuko in v0.5.0 eagerly computed their derived matrices in `init`:
- `RotationTransform`: `self.R = matrix_exp(A - A.mT)`
- `LULinearTransform`: `self.L = tril(LU)`, `self.U = triu(LU, 1) + I`
The cached value detached the transform from its underlying `nn.Parameter`. When a user followed the recommended `F = flow.t()` capture-once pattern (the one `compile_raw` / `compile_beta` documentation centres on), `optimizer.step()` would update the live `LinearMixingTransform.weight` but the captured `F` kept emitting the original `R = I` / `L = I, U = I` from initialisation. The mixing layers degraded to a permanent no-op for any RealNVP trained via the documented compiled path.
Both classes now store the raw input by reference (`self.A` / `self.LU`), expose the derived matrices as `@property` so attribute access recomputes from the live tensor on every read, and override `call_and_ladj` so `ComposedTransform`'s hot path computes the derived matrix exactly once and shares it between `y` and `ladj`. The fix mirrors how `AutoregressiveTransform` already worked via its lazy `meta` closure.
Test coverage added
The bug slipped past v0.5.0 because every existing §14 sub-block (round-trip, `.zeros()`, slogdet, gradient flow, lu-weight-drift) operated on a fresh `flow.t()` after each mutation. Two new sections close the gap:
- §14f (inside the existing `for kind in ("rotation", "lu"):` loop) builds a `RealNVP`, captures `F = flow.t()`, perturbs every parameter in place, and asserts the captured `F.call_and_ladj(x)` matches a freshly built `flow.t().call_and_ladj(x)` at `atol=1e-6`. Pre-fix the rotation path was off by ~1.1 in `y`; post-fix the diff is exactly 0.
- §15 extends the same capture-once consistency check to `NSF` (`MaskedAutoregressiveTransform.meta` closure) and `CNF` (`FreeFormJacobianTransform` module reference). Both pass byte-exact, confirming the lazy-read contract is upheld by every flow class.
Polish (batched in)
- `GeneralCouplingTransform` gains a `.zeros()` method that does the last-MLP-layer weight/bias zeroing previously inlined in `RealNVP.zeros()`. The latter collapses to `for layer in self._layers: layer.zeros()` — same iteration handles couplings and mixing layers uniformly, no more `isinstance` switch.
- `compile_beta` rejects higher-rank `beta` with `assert beta.dim() == 0`. Shape `[1]` or `[N]` would have silently triggered per-shape recompiles, breaking the "one compiled artifact handles every beta" contract. Docstring tightened to "scalar `float | 0-d Tensor`" with an explicit note that per-sample tempering needs a different helper.
Compatibility
The fix is fully backward-compatible for users who never relied on the capture-once contract (i.e. always wrote `reverse_KL(x, target, flow.t())` with a fresh `flow.t()` per call). For users who did follow the documented compiled path with `RealNVP(mixing=...)`, this is a strict correctness improvement: their previously-silent no-op mixing layers will now actually participate in training and may produce noticeably better `reverse_KL` / ESS numbers at
🤖 Generated with Claude Code
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-callas_tensor, lowest-overhead choice for fixed-betatraining.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 ofbeta. Use for annealed / adaptive schedules.- The old
loss.compilename is gone; both helpers live inzflows.lossand 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 mixing=\"rotation\" is an orthogonal mixing=\"lu\" is a PLU map whose learnable None keeps the byte-exact pre-mixing architecture.
Smaller additions
-
Linear_Combinationgeneralised 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-dtorch.Tensor, registered as annn.Modulebuffer 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 = Trueis now the default per-layer feature ordering forNSF/NCSF/RealNVP(freshtorch.randperm(d)each layer);randmask=Falseswitches back to the legacyarange / arange.flipalternation. -
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
Potentialclass with thesuper().__init__()safety warning, install step 4 forcheck_compile_available(), and FAQ entries on conditional flows and JAX alternatives. - New benchmark writeup
tests/compare_compiled_loss.md:compile_rawvs. rawreverse_KLacross adimension × hidden_featuresgrid on an RTX 5070 Ti — 4–10× per-step speedup withmode='default', up to ~13× withmode='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_ladjvs autogradslogdet, NSF/NCSF box preservation, RealNVP forward+inverse ladj cancellation, CNF Hutchinson unbiasedness, the newRealNVPmixing layers, and thecompile_raw/compile_betacache-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
v0.4.0 — drop zuko runtime dependency
Highlights
zflows is now a standalone PyTorch package. The internals previously imported from zuko (Transform subclasses, MaskedMLP, the adaptive ODE solver, the vmap/grad plumbing) have been vendored into zflows.core/ and specialised for the unconditional setting only — every context / c=None argument is gone, so the only distribution a flow ever sees is the one fixed target.
The public API (zflows.flow, zflows.potential, zflows.loss, zflows.utils) and every flow constructor signature is unchanged; downstream code from v0.3.x keeps working with no edits.
Behavioural changes
-
NSF / NCSF parameterise their inner spline on per-coordinate
$[-\text{halfwidth}_i, \text{halfwidth}_i]$ and wrap it with a translation-onlyAdditiveTransformsandwich (no scaling). The box geometry$[a, b]^d$ is honoured natively without distorting the conditioner's dynamic range. -
MonotonicRQSTransformandCircularShiftTransformnow accept a per-coordinateboundtensor, not just a scalar. - The
Flow.t() -> ComposedTransformcontract is preserved across NSF / NCSF / CNF / RealNVP.
Dependencies
- Runtime:
torch>=2.10.0,numpy>=2.4.0. zukois no longer required.
Verification
Two new correctness scripts ship under tests/:
tests/_verify_NSF_NCSF.py— 7 checks (round-trip,zeros()identity,call_and_ladj≈ autogradslogdet, box preservation,CircularShiftper-coord wrap, translation-sandwich invariant).tests/_verify_CNF_RealNVP.py— 7 checks (round-trip,zeros()identity,slogdetagreement, RealNVP forward+inverse ladj cancellation, CNF Hutchinson unbiasedness, backprop coverage).
The end-to-end demos in tests/ all pass against this release.
🤖 Generated with Claude Code