Skip to content

Releases: xuda-ye-math/zflows

v0.6.2

Choose a tag to compare

@xuda-ye-math xuda-ye-math released this 05 Jun 01:44
  • resample: optional N to control number of returned samples
  • docs: add TREE.md with annotated source tree; link from README
  • README: note Potential.enable_grad/enable_eval compile cheaply (safe by default)

v0.6.1: sequential_monte_carlo + hamiltonian_monte_carlo + nested linear_combination

Choose a tag to compare

@xuda-ye-math xuda-ye-math released this 01 Jun 15:39

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

hmchamiltonian_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.py exercises nested combinations end-to-end (depths 2–3, enable cascade, set_coeffs propagation, 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 d with large conditioners (e.g. d=256, hidden_features=(256,256)) compiling the NSF for_ladj/inv_ladj maps is extremely slow to compile, with no effective fix (drop enable_* and use the identical plain maps) — plus a "Why is torch.compile so 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

Choose a tag to compare

@xuda-ye-math xuda-ye-math released this 31 May 23:40

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_Freverse_KL_F, target_KL_Fforward_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.

v0.5.6: lowercase factories + auto-compiled linear_combination

Choose a tag to compare

@xuda-ye-math xuda-ye-math released this 28 May 18:41

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 no lc.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's vmap(grad(forward)) and compile(forward) cost exactly once for the entire $M$-rung anneal.
  • The "needs .enable_grad() first" gate moves from a pre-flight check on Linear_Combination to a runtime check at the child level (an un-enabled child raises RuntimeError the first time the closure tries to call its .grad(x)).

Smaller additions

  • Device-consistency check at Linear_Combination construction. Scans every parameter() and buffer() across the child potentials; raises RuntimeError if 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 previous enable_grad: bool = False / enable_eval: bool = False flags 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 if coeffs is omitted. coeffs is always stored as a plain list[float]; tensor inputs are detached and .tolist()-converted at construction, so there's no buffer-registration / requires_grad subtlety.
  • Docstring caveats on loss_compile / loss_compile_beta. These helpers have only been tested against NSF / NCSF. Behaviour on CNF, OTFlow, RealNVP, or any other flow class is unknown — the docstrings now flag this and recommend verifying against the un-compiled reverse_KL(x, target, flow.t()) baseline before relying on the speedup.
  • Docstring caveat on nesting Linear_Combination. Using a Linear_Combination instance 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 flat linear_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, then linear_combination is 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 uses class MyPotential(Potential) + u = MyPotential().to(device); the factory example uses u = potential_from(myforward).to(device) (the function is myforward, not U_fn); the linear_combination example uses u0, 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 ms mode='default' / 0.5–1.7 ms mode='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 Potential subclasses → potential_from(forward_fn). tests/2D_reverse_KL.py's class U1, tests/_verify_potential.py's class U0 / class U1, tests/_verify_utils.py's class Quartic / class Quartic1D all converted. Stateful subclasses with constructor args — U_target in tests/4D_Boltzmann_generator.py (r0, a, q2, eps), U_target in tests/3D_periodic.py (kappa, renamed from U1), Multi_Well in tests/multi_well_compare.py — kept as classes since potential_from only wraps stateless callables.
  • u0 / u1u_source / u_target across 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.py section B.5 covers the new closure-based behaviour: combined _grad_fn / _eval_fn linked at __init__, runtime gate at the child level, set_coeffs preserves every id() and the closure picks up the new coefficients with err = 0.0 against autograd / forward, enable_grad / enable_eval propagate unconditionally, hot children survive via child-level idempotency, and a linear_combination of pre-enabled children works immediately without lc.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

Choose a tag to compare

@xuda-ye-math xuda-ye-math released this 26 May 14:23

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

Choose a tag to compare

@xuda-ye-math xuda-ye-math released this 26 May 14:14

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

Choose a tag to compare

@xuda-ye-math xuda-ye-math released this 26 May 13:59

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/nvcc and /usr/local/cuda-*/bin/nvcc (Ubuntu's default CUDA-Toolkit install locations, often not on $PATH) as fallbacks when shutil.which(\"nvcc\") misses. Versioned directories are tried newest-first. The warning text enumerates every location searched.
  • GeneralCouplingTransform gained a .zeros() method that does the last-conditioner-MLP weight/bias zero-out previously inlined in RealNVP.zeros(). The latter collapsed to `for layer in self._layers: layer.zeros()` — same iteration handles couplings and mixing layers uniformly.

Correctness fixes

  • LULinearTransform numerical 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 a 1e-12 floor before taking the log — ladj stays finite, gradient saturates to 0 on the affected entry. Healthy LU = I still gives ladj = 0 byte-exact. Docstring updated; the principled Glow-style sign(s0) · exp(log_s) re-parameterisation is noted as a future-major-release option (it would change state_dict layout).
  • compile_beta scalar-only assertion. The wrapper now assert beta.dim() == 0 to 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 "scalar float | 0-d Tensor" with an explicit note that per-sample tempering needs a different helper.
  • Linear_Combination grad-tracking-coeff guard. A user passing `torch.tensor([...], requires_grad=True)` as coeffs would have it registered as a buffer — silently hidden from optimizer.parameters() while backward still accumulated .grad. The constructor now `assert not coeffs.requires_grad` with a clear pointer to coeffs.detach() or the explicit-nn.Parameter escape hatch.

Tests

tests/_verify_flows.py and tests/_verify_potential.py gained four new sub-blocks:

  • §14h — LULinearTransform near-singular safety (diag(L) = 1e-20 still gives finite ladj).
  • §15 — captured-F parameter-mutation consistency on NSF + CNF + OTFlow.
  • §14g — OT_loss = reverse_KL + OT regularizers channel agreement.
  • §B.3 — Linear_Combination rejects requires_grad=True tensor 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 the super().__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_available bullet 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

Choose a tag to compare

@xuda-ye-math xuda-ye-math released this 26 May 04:43

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 $d \gtrsim 8$. No saved checkpoint format changes — `state_dict` keys are identical to v0.5.0.

🤖 Generated with Claude Code

v0.5.0

Choose a tag to compare

@xuda-ye-math xuda-ye-math released this 26 May 04:17

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.pyPotential.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

Choose a tag to compare

@xuda-ye-math xuda-ye-math released this 19 May 00:27

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-only AdditiveTransform sandwich (no scaling). The box geometry $[a, b]^d$ is honoured natively without distorting the conditioner's dynamic range.
  • MonotonicRQSTransform and CircularShiftTransform now accept a per-coordinate bound tensor, not just a scalar.
  • The Flow.t() -> ComposedTransform contract is preserved across NSF / NCSF / CNF / RealNVP.

Dependencies

  • Runtime: torch>=2.10.0, numpy>=2.4.0.
  • zuko is 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 ≈ autograd slogdet, box preservation, CircularShift per-coord wrap, translation-sandwich invariant).
  • tests/_verify_CNF_RealNVP.py — 7 checks (round-trip, zeros() identity, slogdet agreement, 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