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/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