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

🤖 Generated with Claude Code