diff --git a/python/nutpie/compile_pymc.py b/python/nutpie/compile_pymc.py index 13245f7..ca15c8e 100644 --- a/python/nutpie/compile_pymc.py +++ b/python/nutpie/compile_pymc.py @@ -217,6 +217,11 @@ def _make_model(self, init_mean): outer_kwargs = {} def make_adapter(*args, **kwargs): + if "numba_flow" in outer_kwargs: + from nutpie.transform_adapter_numba import make_numba_transform_adapter + + return make_numba_transform_adapter(**outer_kwargs)(*args, **kwargs) + from nutpie.transform_adapter import make_transform_adapter return make_transform_adapter(**outer_kwargs)(*args, **kwargs, logp_fn=None) @@ -233,7 +238,14 @@ def make_adapter(*args, **kwargs): ) def with_transform_adapt(self, **kwargs): - return dataclasses.replace(self, _transform_adapt_args=kwargs) + """Set arguments for the flow transform adapter (``adaptation="flow"``). + + Arguments accumulate across calls; pass ``None`` to reset an + argument to its default. + """ + merged = {**(self._transform_adapt_args or {}), **kwargs} + merged = {k: v for k, v in merged.items() if v is not None} + return dataclasses.replace(self, _transform_adapt_args=merged) def update_user_data(user_data, user_data_storage): @@ -273,6 +285,7 @@ def _compile_pymc_model_numba( model: "pm.Model", pymc_initial_point_fn: Callable[[SeedType], dict[str, np.ndarray]], var_names: Iterable[str] | None = None, + auto_reparam: bool = False, **kwargs, ) -> CompiledPyMCModel: if find_spec("numba") is None: @@ -348,7 +361,7 @@ def _compile_pymc_model_numba( dims, coords = _prepare_dims_and_coords(model, shape_info, reparameterized_names) - return CompiledPyMCModel( + compiled = CompiledPyMCModel( _n_dim=n_dim, dims=dims, _coords=coords, @@ -366,6 +379,16 @@ def _compile_pymc_model_numba( reparameterized_names=reparameterized_names, ) + if auto_reparam: + from nutpie.transform_adapter_numba import build_auto_flow_numba + + # None (with a warning) when the rewrite found nothing to do. + numba_flow = build_auto_flow_numba(model, compiled) + if numba_flow is not None: + compiled = compiled.with_transform_adapt(numba_flow=numba_flow) + + return compiled + def _prepare_dims_and_coords(model, shape_info, reparameterized_names): coords = {} @@ -413,6 +436,7 @@ def _compile_pymc_model_jax( gradient_backend=None, pymc_initial_point_fn: Callable[[SeedType], dict[str, np.ndarray]], var_names: Iterable[str] | None = None, + auto_reparam: bool = False, **kwargs, ): if find_spec("jax") is None: @@ -504,7 +528,7 @@ def expand(_x, **shared): dims, coords = _prepare_dims_and_coords(model, shape_info, reparameterized_names) - return from_pyfunc( + compiled = from_pyfunc( ndim=n_dim, make_logp_fn=make_logp_func, make_expand_fn=make_expand_func, @@ -519,6 +543,16 @@ def expand(_x, **shared): reparameterized_names=reparameterized_names, ) + if auto_reparam: + from nutpie.flow_reparam import build_auto_flow + + # None (with a warning) when the rewrite found nothing to do. + auto_flow = build_auto_flow(model, compiled) + if auto_flow is not None: + compiled = compiled.with_transform_adapt(auto_flow=auto_flow) + + return compiled + def compile_pymc_model( model: "pm.Model", @@ -533,6 +567,7 @@ def compile_pymc_model( ] = "support_point", var_names: Iterable[str] | None = None, freeze_model: bool | None = None, + auto_reparam: bool = False, **kwargs, ) -> CompiledModel: """Compile necessary functions for sampling a pymc model. @@ -561,6 +596,18 @@ def compile_pymc_model( freeze_model : bool | None Freeze all dimensions and shared variables to treat them as compile time constants. + auto_reparam : bool + Automatically reparametrize free random variables (e.g. continuous + VIP centering of location-scale families) and attach the resulting + flow to the model, so that ``nutpie.sample(compiled_model, + adaptation="flow")`` fits the reparametrization during tuning. + Prints a summary of the reparametrized variables; if nothing can be + reparametrized, warns and attaches no flow. Requires + ``backend="jax"`` and ``gradient_backend="jax"``. With a flow + attached, only the reparametrization (plus a diagonal affine) is + fitted by default; further flow options (e.g. ``num_layers`` to add + neural coupling layers) can be set with + ``compiled_model.with_transform_adapt``. Returns ------- @@ -581,6 +628,14 @@ def compile_pymc_model( if backend is not None: backend = backend.lower() # type: ignore[assignment] + # With the jax backend the flow adapter needs the raw JAX logp function, + # which is only kept with the jax gradient backend. The numba backend uses + # the pytensor/numba adapter instead, which builds its own logp graph. + if auto_reparam and backend == "jax" and gradient_backend != "jax": + raise ValueError( + "auto_reparam with backend='jax' requires gradient_backend='jax'" + ) + from pymc.initial_point import make_initial_point_fn from pymc.model.transform.optimization import freeze_dims_and_data @@ -610,6 +665,7 @@ def compile_pymc_model( model=model, pymc_initial_point_fn=initial_point_fn, var_names=var_names, + auto_reparam=auto_reparam, **kwargs, ) elif backend.lower() == "jax": @@ -618,6 +674,7 @@ def compile_pymc_model( gradient_backend=gradient_backend, pymc_initial_point_fn=initial_point_fn, var_names=var_names, + auto_reparam=auto_reparam, **kwargs, ) else: diff --git a/python/nutpie/compiled_pyfunc.py b/python/nutpie/compiled_pyfunc.py index 07602ee..a034ebe 100644 --- a/python/nutpie/compiled_pyfunc.py +++ b/python/nutpie/compiled_pyfunc.py @@ -46,7 +46,15 @@ def with_data(self, **updates): return dataclasses.replace(self, _shared_data=updated) def with_transform_adapt(self, **kwargs): - return dataclasses.replace(self, _transform_adapt_args=kwargs) + """Set arguments for the flow transform adapter (``adaptation="flow"``). + + Arguments accumulate across calls (so e.g. the ``auto_flow`` attached + by ``compile_pymc_model(..., auto_reparam=True)`` survives later + tuning calls); pass ``None`` to reset an argument to its default. + """ + merged = {**(self._transform_adapt_args or {}), **kwargs} + merged = {k: v for k, v in merged.items() if v is not None} + return dataclasses.replace(self, _transform_adapt_args=merged) def _make_sampler( self, diff --git a/python/nutpie/flow_reparam.py b/python/nutpie/flow_reparam.py new file mode 100644 index 0000000..63f54fd --- /dev/null +++ b/python/nutpie/flow_reparam.py @@ -0,0 +1,982 @@ +"""Automatic flow-based reparametrization of free RVs in a PyMC model. + +:func:`reparametrize` picks flows per RV via the extensible rewrite +database :data:`flow_db` and returns plain :class:`FlowSpec` records in +ordinary rv/value/transform space (the rewrite IR never escapes it). +:func:`automatic_flow_reparam` reports the chosen flows, +:func:`build_flow_graph` turns the specs into symbolic +constrain/unconstrain maps over Nutpie's flat point vector, and +:func:`build_auto_flow` wraps those into an ``AutoFlow`` for flow +adaptation. +""" + +from __future__ import annotations + +import abc +import warnings +from typing import NamedTuple + +import numpy as np +import pytensor +import pytensor.tensor as pt +from pytensor.compile import optdb +from pytensor.graph.basic import Apply, clone_get_equiv +from pytensor.graph.fg import FunctionGraph +from pytensor.graph.rewriting.basic import dfs_rewriter, node_rewriter +from pytensor.graph.rewriting.db import RewriteDatabaseQuery, SequenceDB +from pytensor.graph.traversal import ancestors, explicit_graph_inputs +from pytensor.tensor.math import variadic_add +from pytensor.tensor.random.basic import ( + CauchyRV, + ExponentialRV, + GammaRV, + GumbelRV, + HalfCauchyRV, + HalfNormalRV, + InvGammaRV, + LaplaceRV, + LogisticRV, + LogNormalRV, + NormalRV, + ParetoRV, + StudentTRV, + WeibullRV, +) +from pytensor.xtensor.basic import ( + XTensorFromTensor, + tensor_from_xtensor, + xtensor_from_tensor, +) +from pytensor.xtensor.type import XTensorType + + +from pymc.dims.distributions.transforms import LogTransform as DimLogTransform +from pymc.distributions.multivariate import ZeroSumNormalRV +from pymc.distributions.transforms import ZeroSumTransform +from pymc.distributions.transforms import log as log_transform +from pymc.logprob.transforms import LogTransform +from pymc.logprob.utils import replace_rvs_by_values +from pymc.model.core import Model +from pymc.model.fgraph import ( + ModelFreeRV, + ModelValuedVar, + ModelVar, + fgraph_from_model, +) +from pymc.pytensorf import toposort_replace + + +# (loc_idx, scale_idx, required_value_transform) per RV op type. The transform +# entry is the class the ModelFreeRV's value transform must be for this RV +# op's sampling-space distribution to be loc-scale. ``None`` means no +# transform (rv and value share the same space, already loc-scale). +# ``LogTransform`` covers log-loc-scale families (LogNormal in log-space +# is Normal, so the same affine flow applies). +LOC_SCALE_FAMILIES: dict[type, tuple[int, int, type | None]] = { + NormalRV: (0, 1, None), + CauchyRV: (0, 1, None), + LaplaceRV: (0, 1, None), + LogisticRV: (0, 1, None), + GumbelRV: (0, 1, None), + StudentTRV: (1, 2, None), + LogNormalRV: (0, 1, LogTransform), +} + + +# RV ops of the form ``X = S · Y`` (``Y`` fixed-shape, ``S`` the scale) +# that land in log-space under their default ``LogTransform``. Shifting +# the log-space value by a constant ↔ scaling ``X`` by that constant, so +# a shift-only affine flow in the sampling space captures hierarchical +# variation in the scale parameter without an icdf call. Value is the +# index of the scale/rate parameter within ``dist_params``. +SCALE_SHIFT_FAMILIES: dict[type, int] = { + GammaRV: 1, # β (rate) + WeibullRV: 1, # β + InvGammaRV: 1, # β + ExponentialRV: 0, # λ + HalfNormalRV: 0, # σ + HalfCauchyRV: 0, # β + ParetoRV: 1, # m (scale; α fixed) +} + + +class Flow(abc.ABC): + """Pure-math descriptor for a per-RV reparametrization. + + Subclasses declare how many of the ``constrain`` / ``unconstrain`` + arguments come from the model graph (``n_model_params`` — e.g. the + location/scale a child reads off its parents) versus how many are + trainable per-flow parameters fitted during adaptation + (``n_hyper_params`` — the VIP centering knobs). Both maps take the + same signature ``(point, *model_params, *hyper_params)``: + + - ``unconstrain(x, ...)`` — value-space point ``x`` → NUTS-space + point ``y``. + - ``constrain(y, ...)`` — NUTS-space point ``y`` → value-space + point ``x``. + + Both log|det| Jacobians default to the full Jacobian via + ``pt.jacobian(..., vectorize=True)`` — correct for any invertible + flow; subclasses override with a closed form when it's cheaper (see + :class:`BaseAffineFlow` for the constant-Jacobian shortcut). + """ + + # Leading args sourced from the model fgraph (loc/scale subgraphs). + n_model_params: int = 0 + # Trailing args fitted during adaptation; one entry per ``param_shapes``. + n_hyper_params: int = 0 + + @staticmethod + @abc.abstractmethod + def unconstrain(x, *params): # x -> y + ... + + @staticmethod + @abc.abstractmethod + def constrain(y, *params): # y -> x + ... + + @classmethod + def log_jac_det_constrain(cls, y, *params): + x = cls.constrain(y, *params).ravel() + J = pt.jacobian(x, y, vectorize=True).reshape((x.size, x.size)) + return pt.linalg.slogdet(J)[1] + + @classmethod + def log_jac_det_unconstrain(cls, x, *params): + y = cls.unconstrain(x, *params).ravel() + J = pt.jacobian(y, x, vectorize=True).reshape((y.size, y.size)) + return pt.linalg.slogdet(J)[1] + + +class BaseAffineFlow(Flow): + """Flow that is affine in the point, i.e. its Jacobian is constant. + + For such flows the ``unconstrain`` log|det| is just the negation of + the ``constrain`` one, and the point at which the formula is + evaluated is irrelevant — so only ``log_jac_det_constrain`` needs a + closed form. + """ + + @classmethod + def log_jac_det_unconstrain(cls, x, *params): + return -cls.log_jac_det_constrain(x, *params) + + +class NoFlow(BaseAffineFlow): + """Identity flow. Used for RVs the rewrite skipped.""" + + n_model_params = 0 + n_hyper_params = 0 + + @staticmethod + def unconstrain(x): + return x + + @staticmethod + def constrain(y): + return y + + @staticmethod + def log_jac_det_constrain(y): + return pt.zeros((), dtype=y.dtype) + + +def _pin_if_empty(h): + """Hyper params arrive with concrete static shapes (they are unpacked + from the flat trainable vector); size 0 marks a knob the rewrite + withheld because its dist param did not qualify. Pin it at the + centred no-op ``h = 0`` so it drops out of the transform.""" + if 0 in h.type.shape: + return pt.zeros((), dtype=h.dtype) + return h + + +class AffineFlow(BaseAffineFlow): + """Variationally Inferred Parameterisation of a location-scale RV. + + Following Gorinova et al. (2019), a child ``z ~ Dist(loc, scale)`` is + expressed via a standardized ``y`` interpolating continuously between + the centred (CP) and non-centred (NCP) parameterisations. A single + trainable knob ``h`` per variable (the paper's ``1 - λ``) controls + both location and scale, so ``h = 0`` is the no-op (centred):: + + constrain(y, loc, scale, h) = (y - (1 - h)·loc)·scale^h + loc + + ``h = 0`` ⇒ identity (CP); ``h = 1`` ⇒ full NCP (``y`` is the + standardized residual). ``loc`` and ``scale`` are read off the parent + RVs (model params); ``h`` is the trainable hyper param, left + unconstrained (over-/under-centering is allowed). + + The hyper params broadcast against ``y`` / ``x`` via normal numpy + rules. The rewrite allocates one knob per element of the value var + (following Gorinova et al., whose λ is shaped like the RV), so + different elements of one hierarchical group can settle on different + centerings — except on axes where elements cannot be reparametrized + independently (e.g. ZeroSumNormal core dims), which get size 1. + + A dist param can also fail to qualify for a knob altogether (constant + loc, full-shape scale, ...). The rewrite then allocates its hyper + param at size 0, and the flow pins that knob at the centred no-op + ``h = 0`` (see :func:`_pin_if_empty`) — e.g. with ``h_sigma`` pinned + the transform degenerates to the translation ``y + h_mu·loc``. + """ + + n_model_params = 2 + n_hyper_params = 2 + + @staticmethod + def unconstrain(x, pop_mu, pop_sigma, h_mu, h_sigma): + h_mu, h_sigma = _pin_if_empty(h_mu), _pin_if_empty(h_sigma) + return (x - pop_mu) * (pop_sigma**-h_sigma) + (1 - h_mu) * pop_mu + + @staticmethod + def constrain(y, pop_mu, pop_sigma, h_mu, h_sigma): + h_mu, h_sigma = _pin_if_empty(h_mu), _pin_if_empty(h_sigma) + return (y - (1 - h_mu) * pop_mu) * (pop_sigma**h_sigma) + pop_mu + + @staticmethod + def log_jac_det_constrain(y, pop_mu, pop_sigma, h_mu, h_sigma): + # d(value)/d(y) = scale^h; h = 0 ⇒ identity. + log_det = _pin_if_empty(h_sigma) * pt.log(pop_sigma) + return pt.broadcast_to(log_det, y.shape).sum() + + +class ShiftFlow(BaseAffineFlow): + """VIP reparametrization of a log-transformed scale-family RV. + + For ``X`` whose default ``LogTransform`` lands it in log-space, the + parent enters the log-space value additively through ``log(scale)``, + so non-centering is a *shift* by that amount (no scale exponent):: + + constrain(y, shift, h) = y + h·shift with shift = log(scale) + + ``h = 0`` is the no-op (centred); ``h`` is unconstrained, so the + optimizer reaches full decoupling regardless of whether the model + param is a rate or a scale (the sign is absorbed into ``h``). The + Jacobian is a translation (det = 1 ⇒ log|det| = 0). + """ + + n_model_params = 1 + n_hyper_params = 1 + + @staticmethod + def unconstrain(x, shift, h): + return x - h * shift + + @staticmethod + def constrain(y, shift, h): + return y + h * shift + + @staticmethod + def log_jac_det_constrain(y, shift, h): + return pt.zeros((), dtype=y.dtype) + + +class FlowSpec(NamedTuple): + """Per-free-RV reparametrization record in plain rv/value/transform + space — the rewrite IR never escapes :func:`reparametrize`. + + ``rv`` is the genuine random variable (so PyMC's + ``replace_rvs_by_values`` passes the actual distribution parameters + to ``transform.backward``); ``model_params`` are the loc/scale + subgraphs the flow reads off the model, expressed in terms of the + parent specs' ``rv`` variables; ``param_shapes`` are the concrete + shapes of the flow's trainable hyper params, evaluated at the model's + initial point. + """ + + flow_cls: type[Flow] + rv: object + value: object + transform: object | None + model_params: list + param_shapes: list[tuple[int, ...]] + + +class FlowFreeRV(ModelValuedVar): + """Rewrite-internal marker carrying the chosen flow, the RV's value + transform, the model-derived params, and per-hyper-param shape exprs. + Lives only between the flow rewrite and the IR strip inside + :func:`reparametrize`. + + Inputs (positional, flat so ``op.make_node(*node.inputs)`` is + idempotent):: + + (rv, value, *model_params, *hyper_shape_exprs) + + with ``len(model_params) == flow_cls.n_model_params`` and + ``len(hyper_shape_exprs) == flow_cls.n_hyper_params``. + """ + + __props__ = ("name", "dims", "transform", "flow_cls") + + def __init__(self, flow_cls: type[Flow], name, dims=(), transform=None): + self.flow_cls = flow_cls + super().__init__(name, dims, transform=transform) + + def __call__(self, rv, value, *params, hyperparam_shapes=()): + # Ergonomic construction: callers group the variadic chunks by name. + return super().__call__(rv, value, *params, *hyperparam_shapes) + + def make_node(self, rv, value, *rest): + nm = self.flow_cls.n_model_params + nh = self.flow_cls.n_hyper_params + assert len(rest) >= nm + nh + return Apply(self, [rv, value, *rest], [value.type(name=value.name)]) + + +class _FlowParts(NamedTuple): + """IR-side view of a (Flow|Model)FreeRV node, used only while the + rewritten fgraph is alive.""" + + flow_cls: type[Flow] + rv: object + value: object + out: object + transform: object | None + model_params: list + hyper_shapes: list + + +def _flow_node_parts(node) -> _FlowParts | None: + """Extract :class:`_FlowParts` from a ``FlowFreeRV`` or plain + ``ModelFreeRV`` node (the latter mapped to :class:`NoFlow`); returns + ``None`` for any other node.""" + op = node.op + if isinstance(op, FlowFreeRV): + rv, value, *rest = node.inputs + nm = op.flow_cls.n_model_params + nh = op.flow_cls.n_hyper_params + return _FlowParts( + flow_cls=op.flow_cls, + rv=rv, + value=value, + out=node.outputs[0], + transform=op.transform, + model_params=list(rest[:nm]), + hyper_shapes=list(rest[nm : nm + nh]), + ) + if isinstance(op, ModelFreeRV): + rv, value, *_dims = node.inputs + return _FlowParts( + flow_cls=NoFlow, + rv=rv, + value=value, + out=node.outputs[0], + transform=op.transform, + model_params=[], + hyper_shapes=[], + ) + return None + + +def _depends_on_free_rv(vars_) -> bool: + return any( + anc.owner is not None and isinstance(anc.owner.op, ModelFreeRV) + for anc in ancestors(vars_) + ) + + +@node_rewriter([ModelFreeRV]) +def lift_xtensor_from_model_free_rv(fgraph, node): + """Pull ``XTensorFromTensor`` wrappers out of a ``ModelFreeRV``'s rv + and value inputs, leaving plain tensors inside so downstream flow + rewrites don't have to know about xtensor:: + + ModelFreeRV(XTensorFromTensor(rv), XTensorFromTensor(value)) + -> XTensorFromTensor(ModelFreeRV(rv, aligned_value)) + + Only fires on RVs with ``transform=None`` or a :class:`DimTransform` + known to have a plain counterpart (see the ``match`` below); unknown + transforms are left alone since peeling might change semantics — + such RVs stay xtensor-typed end to end, which the spec extraction + and graph build handle natively. Invariant: rv and value are both + xtensor (or both plain tensors) — they may declare different dim + orders, in which case the value is ``dimshuffle``d into rv's order + so the rebuilt inner ``ModelFreeRV`` sees matched-axis tensors. + """ + current_transform = node.op.transform + # Swap any dim-aware transform for its plain logprob counterpart so + # downstream rewrites see a single class hierarchy. Unknown transforms + # aren't safe to peel — leave the node alone. + match current_transform: + case None: + new_transform = None + case DimLogTransform(): + new_transform = log_transform + case _: + return None + + xrv, xvalue = node.inputs + if not isinstance(xrv.owner.op, XTensorFromTensor): + return None + rv_dims = xrv.type.dims + value_dims = xvalue.type.dims + if set(rv_dims) != set(value_dims): + # Unrelated named axes — no safe permutation; leave the node alone. + return None + rv = xrv.owner.inputs[0] + if xvalue.owner is not None and isinstance(xvalue.owner.op, XTensorFromTensor): + value = xvalue.owner.inputs[0] + else: + value = tensor_from_xtensor(xvalue) + value.name = xvalue.name + if rv_dims != value_dims: + value = value.dimshuffle([value_dims.index(d) for d in rv_dims]) + value.name = xvalue.name + new_op = ( + node.op + if new_transform is current_transform + else type(node.op)(node.op.name, node.op.dims, transform=new_transform) + ) + new_free_rv = new_op(rv, value) + return [XTensorFromTensor(dims=rv_dims)(new_free_rv)] + + +_one = pt.constant(1) +_empty_shape = pt.constant(np.array([0], dtype="int64")) + + +def _stack_shape(dims): + """Stack scalar shape components into a 1-D int64 shape vector, + returning the empty vector (shape ``()``) when there are none — plain + ``pt.stack([])`` raises ``"No tensor arguments provided."``.""" + if not dims: + return pt.constant(np.array([], dtype="int64")) + return pt.stack(dims) + + +def _hyper_shape(value, n_core: int = 0): + """Per-element hyper-param shape: the value var's shape, collapsed to + size 1 on the trailing ``n_core`` axes (elements along support dims — + e.g. ZeroSumNormal core dims — cannot be reparametrized independently + and must share one knob).""" + ndim = value.ndim + dims = [*(value.shape[i] for i in range(ndim - n_core)), *([_one] * n_core)] + return _stack_shape(dims) + + +@node_rewriter([ModelFreeRV]) +def loc_scale_affine_flow(fgraph, node): + rv, value = node.inputs + rv_node = rv.owner + + entry = LOC_SCALE_FAMILIES.get(type(rv_node.op)) + if entry is None: + return None + + loc_idx, scale_idx, expected_transform = entry + if expected_transform is None: + if node.op.transform is not None: + return None + elif not isinstance(node.op.transform, expected_transform): + return None + + dist_params = list(rv_node.op.dist_params(rv_node)) + loc, scale = dist_params[loc_idx], dist_params[scale_idx] + # A param earns a centering knob when it carries free-RV randomness — + # i.e. the modeller left the RV (fully or partly) centred on another + # latent, so there is prior dependence to decouple. Where centring is + # already best the fitted knob collapses to the no-op (h = 0). + loc_qualifies = _depends_on_free_rv([loc]) + scale_qualifies = _depends_on_free_rv([scale]) + if not (loc_qualifies or scale_qualifies): + return None + + # Each dist param earns its own per-element knob; a param that does + # not qualify gets a size-0 hyper param, which the flow pins at the + # centred no-op. + param_shape = _hyper_shape(value) + flow_rv = FlowFreeRV(AffineFlow, **node.op._props_dict())( + rv, + value, + loc, + scale, + hyperparam_shapes=[ + param_shape if loc_qualifies else _empty_shape, + param_shape if scale_qualifies else _empty_shape, + ], + ) + return {node.outputs[0]: flow_rv} + + +@node_rewriter([ModelFreeRV]) +def scale_shift_flow(fgraph, node): + if not isinstance(node.op.transform, LogTransform): + return None + + rv, value = node.inputs + rv_node = rv.owner + + scale_idx = SCALE_SHIFT_FAMILIES.get(type(rv_node.op)) + if scale_idx is None: + return None + + scale = rv_node.op.dist_params(rv_node)[scale_idx] + if not _depends_on_free_rv([scale]): + return None + + # VIP shift: the parent enters the log-space value additively through + # ``log(scale)``, so the flow shifts by ``h·log(scale)``; h = 0 is + # centred. + flow_rv = FlowFreeRV(ShiftFlow, **node.op._props_dict())( + rv, value, pt.log(scale), hyperparam_shapes=[_hyper_shape(value)] + ) + return {node.outputs[0]: flow_rv} + + +@node_rewriter([ModelFreeRV]) +def zerosum_scale_flow(fgraph, node): + """VIP scale flow for ``ZeroSumNormal`` with a common (per-batch) + sigma. + + Under the default ``ZeroSumTransform`` the value var is iid + ``Normal(0, σ)`` over the reduced support dims (the transform is an + isometry onto the zero-sum hyperplane and σ has core shape 1, so the + scaling commutes with it). That makes the value-space RV a zero-loc + scale family: :class:`AffineFlow` with ``loc = 0`` (loc knob + withheld) applies, with one scale knob per batch element shared + across the core dims (a per-element knob there would break the + zero-sum coupling). + """ + if not isinstance(node.op.transform, ZeroSumTransform): + return None + + rv, value = node.inputs + rv_node = rv.owner + if not isinstance(rv_node.op, ZeroSumNormalRV): + return None + + # ZeroSumNormalRV constructs sigma with core shape (1, ...), so it is + # always broadcast along the core dims; guard the invariant anyway. + sigma = rv_node.op.dist_params(rv_node)[0] + if not all(sigma.type.broadcastable[sigma.type.ndim - rv_node.op.ndim_supp :]): + return None + if not _depends_on_free_rv([sigma]): + return None + + n_core = rv_node.op.ndim_supp + flow_rv = FlowFreeRV(AffineFlow, **node.op._props_dict())( + rv, + value, + pt.zeros((), dtype=value.dtype), + sigma, + hyperparam_shapes=[_empty_shape, _hyper_shape(value, n_core=n_core)], + ) + return {node.outputs[0]: flow_rv} + + +# Tag taxonomy: +# "default" — plumbing + safe-by-default flow rewrites (used by the default +# query); always produces correct posteriors. +# "all" — everything, including opt-in rewrites that a user may want +# finer-grained control over. +# Per-flow tags ("affine", "icdf") allow targeted selection. +default_flow_query = RewriteDatabaseQuery(include=("default",)) +flow_db = SequenceDB() +flow_db.register("lower_xtensor", optdb.query("+lower_xtensor"), "default", "all") +flow_db.register( + "lift_xtensor_from_model_free_rv", + dfs_rewriter(lift_xtensor_from_model_free_rv), + "default", + "all", +) +flow_db.register( + "affine_flow", + dfs_rewriter( + loc_scale_affine_flow, + scale_shift_flow, + zerosum_scale_flow, + ), + "default", + "all", + "affine", +) + + +def _eval_hyper_shapes( + model: Model, parts: list[_FlowParts] +) -> dict[str, list[tuple[int, ...]]]: + """Evaluate every flow's symbolic hyper-param shape expressions at the + model's initial point, in one compile. Keyed by value-var name.""" + exprs = [s for p in parts for s in p.hyper_shapes] + if exprs: + inputs = list(explicit_graph_inputs(exprs)) + ip = model.initial_point() + vals = pytensor.function(inputs, exprs)(**{v.name: ip[v.name] for v in inputs}) + else: + vals = [] + shapes: dict[str, list[tuple[int, ...]]] = {} + idx = 0 + for p in parts: + k = p.flow_cls.n_hyper_params + shapes[p.value.name] = [tuple(int(x) for x in s) for s in vals[idx : idx + k]] + idx += k + return shapes + + +def _first_non_model_var(var): + while var.owner is not None and isinstance(var.owner.op, ModelVar): + op = var.owner.op + # Observed RVs (valued, but not free/flow) resolve to their observed + # value (data, inputs[1]); using the rv (inputs[0]) would leak its + # RNG into the flow graph (PyMC's logp likewise substitutes observed + # RVs by their data). Free/flow RVs and other ModelVars use inputs[0]. + if isinstance(op, ModelValuedVar) and not isinstance( + op, (ModelFreeRV, FlowFreeRV) + ): + var = var.owner.inputs[1] + else: + var = var.owner.inputs[0] + return var + + +def reparametrize( + model: Model, + flow_db_query: RewriteDatabaseQuery = default_flow_query, + db: SequenceDB = flow_db, +) -> list[FlowSpec]: + """Run the flow rewrite and return one :class:`FlowSpec` per free RV, + in fgraph toposort order. + + The rewrite happens on the PyMC model IR (``fgraph_from_model``); + afterwards every ``ModelVar`` dummy is stripped — the same in-place + replacement ``model_from_fgraph`` performs — so the returned specs + live in plain rv/value/transform space. PyMC's + ``replace_rvs_by_values`` then composes parent dependencies and value + transforms correctly on them, including conditional transforms (which + read the RV's actual distribution parameters) and xtensor variables. + """ + fgraph, _memo = fgraph_from_model(model) + db.query(flow_db_query).rewrite(fgraph) + # Per-free-RV parts, in fgraph toposort order. + parts = [ + p for node in fgraph.toposort() if (p := _flow_node_parts(node)) is not None + ] + shapes = _eval_hyper_shapes(model, parts) + # Resolve held references to non-dummy vars *before* the strip: the + # in-place replacement below rewires the ancestors of vars that stay + # in the fgraph, but vars removed from it keep stale inputs. + resolved = [ + (p, [_first_non_model_var(g) for g in (p.rv, *p.model_params)]) for p in parts + ] + # Strip the IR in place (cf. ``model_from_fgraph``). Forward toposort + # order: a dummy is always replaced while its consumers are still in + # the graph, so subgraphs only reachable through a later-stripped + # FlowFreeRV node are rewired before they are pruned. + dummy_replacements = [] + for node in fgraph.toposort(): + op = node.op + if not isinstance(op, ModelVar): + continue + # Observed RVs strip to their observed value (data, inputs[1]); the rv + # (inputs[0]) carries an RNG that would leak into the flow graph. Free + # RVs become their rv (replaced by value vars later); other ModelVars + # (deterministics, …) use inputs[0]. + if isinstance(op, ModelValuedVar) and not isinstance( + op, (ModelFreeRV, FlowFreeRV) + ): + repl = _first_non_model_var(node.inputs[1]) + else: + repl = _first_non_model_var(node.inputs[0]) + dummy_replacements.append((node.outputs[0], repl)) + toposort_replace(fgraph, dummy_replacements) + return [ + FlowSpec( + flow_cls=p.flow_cls, + rv=rv, + value=p.value, + transform=p.transform, + model_params=model_params, + param_shapes=shapes[p.value.name], + ) + for p, (rv, *model_params) in resolved + ] + + +def automatic_flow_reparam( + model: Model, + flow_db_query: RewriteDatabaseQuery = default_flow_query, + db: SequenceDB = flow_db, +) -> dict[str, dict]: + """Run the flow rewrite and report, per unconstrained value variable, + which flow was chosen and the concrete shapes of its hyper params. + + Returns a ``dict`` keyed by the value variable's name (insertion order + matches fgraph toposort). Each value is a dict with: + + ``flow_cls`` — the :class:`Flow` descriptor class (``NoFlow`` when the + rewrite skipped the RV). + ``dtype`` — dtype of the flow's hyper params. + ``transform`` — the RV's value transform (``None`` or a PyMC + :class:`Transform`). + ``param_shapes`` — ``list[tuple[int, ...]]`` concrete shape per hyper + param. + """ + specs = reparametrize(model, flow_db_query, db) + return { + s.value.name: dict( + flow_cls=s.flow_cls, + dtype=s.rv.type.dtype, + transform=s.transform, + param_shapes=s.param_shapes, + ) + for s in specs + } + + +def build_flow_graph_from_specs( + specs: list[FlowSpec], + free_vars_info, + n_dim: int, +) -> dict[str, object]: + """Build the symbolic constrain/unconstrain flow maps over Nutpie's + flat point vector, plus a flat trainable flow-params vector. + + Each RV's ``constrain`` / ``unconstrain`` is first expressed against + fresh value/hyper-param placeholders with its value transform folded + in, then PyMC's :func:`replace_rvs_by_values` composes the + parent→child dependency in topological order. A final + ``toposort_replace`` substitutes each value placeholder by its chunk + of the flat point vector and each hyper placeholder by its slice of + the flat params vector — since those are introduced last, the + returned inputs are exact (no recovering cloned inputs by name). + + Parameters + ---------- + specs + Output of :func:`reparametrize`. + free_vars_info + Per-free-variable descriptors whose ``.name``, ``.start_idx``, + ``.end_idx`` and ``.shape`` define the flat point vector layout — + typically ``compiled_model._variables`` filtered to the free + (unconstrained) ones. Order determines packing into the flat + vector and must match Nutpie's. + n_dim + Total dimension of the flat point vector (``compiled_model.n_dim``). + + Returns + ------- + dict with keys: + ``flow_params_vector`` — flat trainable params ``pt.vector``. + ``constrain`` — ``(inputs, outputs)`` for ``y -> value``: + inputs ``[y_vector, flow_params_vector]``, + outputs ``[value_point, total_log_jac_det_constrain]``. + ``unconstrain`` — ``(inputs, outputs)`` for ``value -> y``. + """ + n_dim = int(n_dim) + order = [v.name for v in free_vars_info] + info = {v.name: v for v in free_vars_info} + + # Flat trainable params vector, sliced into each flow's hyper params. + param_shapes = [sh for s in specs for sh in s.param_shapes] + total = int(sum(np.prod(sh) for sh in param_shapes)) if param_shapes else 0 + flow_params = pt.vector("flow_params", dtype="float64", shape=(total,)) + splits = list(pt.unpack(flow_params, param_shapes)) if param_shapes else [] + # Fresh placeholder per hyper param; substituted by its split last. + hyper: dict[str, list] = {} + hyper_to_split: dict = {} + idx = 0 + for s in specs: + phs = [ + pt.tensor(f"{s.value.name}_hyper{i}", shape=sh, dtype="float64") + for i, sh in enumerate(s.param_shapes) + ] + hyper[s.value.name] = phs + for ph in phs: + hyper_to_split[ph] = splits[idx] + idx += 1 + + def _root_and_chunk(vec, spec): + # The lift rewrite leaves lifted values as derived expressions + # (``tensor_from_xtensor(xvalue)``); the flat chunk substitutes + # the *root* value var, in its own layout. + if spec.value.owner is None: + root = spec.value + else: + (root,) = explicit_graph_inputs([spec.value]) + v = info[spec.value.name] + chunk = vec[v.start_idx : v.end_idx].reshape(tuple(int(x) for x in v.shape)) + if isinstance(root.type, XTensorType): + chunk = xtensor_from_tensor(chunk, dims=root.type.dims, name=root.name) + return root, chunk + + def _build(direction: str): + y = pt.vector("y", shape=(n_dim,)) + # Per-direction copy of the model-param subgraphs: + # replace_rvs_by_values mutates replacement expressions in place + # when they nest other replaced rvs (see replace_vars_in_graphs), + # so the shared spec graphs must not be fed to it directly. The + # rv keys are pinned to identity so they stay valid keys. + memo = {s.rv: s.rv for s in specs} + all_params = [p for s in specs for p in s.model_params] + equiv = clone_get_equiv([], all_params, False, False, memo) + model_params = {s.value.name: [equiv[p] for p in s.model_params] for s in specs} + points: dict[str, object] = {} + ljds: dict[str, object] = {} + rvs_to_values: dict = {} + replacements: list = [] + for s in specs: + name = s.value.name + # Fresh root placeholder for this RV's flat chunk, substituted + # at the end; cloning inside replace_rvs_by_values keeps graph + # inputs identical, so the substitution is exact. The value + # var's derivation (if any) is rebuilt on top of it so the + # flow math sees the same layout as the spec graphs. + root, chunk = _root_and_chunk(y, s) + z_root = root.type(name=root.name) + replacements.append((z_root, chunk)) + if s.value is root: + z = z_root + else: + memo_v = clone_get_equiv( + [root], [s.value], False, False, {root: z_root} + ) + z = memo_v[s.value] + params = model_params[name] + if s.flow_cls is NoFlow: + point, ljd = z, pt.zeros(()) + elif direction == "constrain": + point = s.flow_cls.constrain(z, *params, *hyper[name]) + ljd = s.flow_cls.log_jac_det_constrain(z, *params, *hyper[name]) + else: + point = s.flow_cls.unconstrain(z, *params, *hyper[name]) + ljd = s.flow_cls.log_jac_det_unconstrain(z, *params, *hyper[name]) + points[name], ljds[name] = point, ljd + # A child reads this RV's *constrained* value off its parents: + # the flow output (value space) in constrain, the value var in + # unconstrain — backward-transformed with the RV's actual + # distribution parameters, so conditional transforms compose + # correctly. The backward is folded into the value here + # instead of passing rvs_to_transforms: with transforms, + # replace_rvs_by_values clones the graphs and remaps the + # *keys* to clones, so rvs nested inside other replacement + # values (a flow parent's point expression) would be missed. + rv_value = point if direction == "constrain" else z + if s.transform is not None: + rv_value = s.transform.backward(rv_value, *s.rv.owner.inputs) + rv_value = s.rv.type.filter_variable(rv_value, allow_convert=True) + rv_value.name = s.rv.name + rvs_to_values[s.rv] = rv_value + + graphs = [points[nm] for nm in order] + [ljds[nm] for nm in order] + graphs = replace_rvs_by_values(graphs, rvs_to_values=rvs_to_values) + n = len(order) + point_parts = [ + tensor_from_xtensor(g) if isinstance(g.type, XTensorType) else g + for g in graphs[:n] + ] + point_out = pt.concatenate([g.ravel() for g in point_parts]) + ljd_out = variadic_add(*graphs[n:]) + # Substitute the value/hyper placeholders by their flat-vector + # slices. replace_rvs_by_values keeps graph inputs identical when + # cloning, so the placeholders (and thus these replacements) are + # exact. + fg = FunctionGraph(outputs=[point_out, ljd_out], clone=False) + final_replacements = [ + (ph, repl) + for ph, repl in (*replacements, *hyper_to_split.items()) + if ph in fg.variables + ] + toposort_replace(fg, final_replacements) + return [y, flow_params], list(fg.outputs) + + return dict( + flow_params_vector=flow_params, + constrain=_build("constrain"), + unconstrain=_build("unconstrain"), + ) + + +def build_flow_graph( + model: Model, + free_vars_info, + n_dim: int, + flow_db_query: RewriteDatabaseQuery = default_flow_query, + db: SequenceDB = flow_db, +) -> dict[str, object]: + """:func:`reparametrize` + :func:`build_flow_graph_from_specs`.""" + specs = reparametrize(model, flow_db_query, db) + return build_flow_graph_from_specs(specs, free_vars_info, n_dim) + + +class VarInfo(NamedTuple): + name: str + start_idx: int + end_idx: int + shape: tuple[int, ...] + + +def free_vars_info(compiled_model): + """The compiled model's free (unconstrained) variable descriptors, + whose ``start_idx``/``end_idx``/``shape`` define Nutpie's flat point + vector layout.""" + n_dim = int(compiled_model.n_dim) + if hasattr(compiled_model, "_variables"): # pyfunc (jax) model + return [v for v in compiled_model._variables if v.end_idx <= n_dim] + # numba model: the free vars are the leading entries of the flat vector + names, slices, shapes = compiled_model.shape_info + return [ + VarInfo(name, sl.start, sl.stop, tuple(shape)) + for name, sl, shape in zip(names, slices, shapes, strict=True) + if sl.stop <= n_dim + ] + + +def build_auto_flow( + model: Model, + compiled_model, + *, + init_params=None, + flow_db_query: RewriteDatabaseQuery = default_flow_query, + db: SequenceDB = flow_db, +): + """Build the VIP reparametrization of ``model`` as a single + :class:`nutpie.normalizing_flow.AutoFlow` over ``compiled_model``'s + flat point vector, ready to pass to + ``compiled_model.with_transform_adapt(auto_flow=...)``. + + Prints a summary of the reparametrized variables; if the rewrite found + nothing to reparametrize, warns and returns ``None`` instead (use + :func:`automatic_flow_reparam` for the full per-variable report). + + The PyTensor constrain/unconstrain maps from :func:`build_flow_graph` + are JIT-compiled to JAX. The flow's trainable parameters are the VIP + centering knobs ``h`` (the paper's ``1 - λ``); ``init_params`` defaults + to zeros, i.e. ``h = 0`` — the centred (CP) no-op, so adaptation starts + from the original model and moves toward non-centring as needed (see + :class:`~nutpie.flow_reparam.AffineFlow`). + + The flowjax/JAX bijection convention is ``transform: base -> target``, + matching Nutpie's ``transform_and_log_det(sampler) -> value``; so the + flow's ``transform`` is :func:`build_flow_graph`'s ``constrain`` and + its ``inverse`` is ``unconstrain``. + """ + from nutpie.normalizing_flow import AutoFlow + import jax.numpy as jnp + + specs = reparametrize(model, flow_db_query, db) + flowed = [s for s in specs if s.flow_cls is not NoFlow] + if not flowed: + warnings.warn( + "Automatic reparametrization did not find any variables to " + "reparametrize in this model." + ) + return None + chosen = ", ".join(f"{s.value.name} ({s.flow_cls.__name__})" for s in flowed) + print( + f"auto_reparam: reparametrizing {len(flowed)} of {len(specs)} " + f"free variables: {chosen}" + ) + + n_dim = int(compiled_model.n_dim) + g = build_flow_graph_from_specs(specs, free_vars_info(compiled_model), n_dim) + constrain_fn = pytensor.function(*g["constrain"], mode="JAX").vm.jit_fn + unconstrain_fn = pytensor.function(*g["unconstrain"], mode="JAX").vm.jit_fn + + total = int(g["flow_params_vector"].type.shape[0]) + if init_params is None: + init_params = jnp.zeros((total,)) + + return AutoFlow(init_params, (n_dim,), constrain_fn, unconstrain_fn) diff --git a/python/nutpie/normalizing_flow.py b/python/nutpie/normalizing_flow.py index 97c54be..7909fee 100644 --- a/python/nutpie/normalizing_flow.py +++ b/python/nutpie/normalizing_flow.py @@ -18,6 +18,17 @@ from paramax.wrappers import AbstractUnwrappable +def _positive(x): + """Unconstrained -> positive reparametrization (``exp(asinh(x))``). + + A module-level function, not a lambda: ``Parameterize`` keeps it as a + static pytree leaf, so a fresh lambda per flow would change the static + structure and force jax to retrace and recompile every jitted function + touching the flow on every adaptation window. + """ + return x + jnp.sqrt(1 + x**2) + + def _generate_sequences(k, r_vals): """ Generate all binary sequences of length k with exactly r 1's. @@ -342,8 +353,8 @@ def __init__( ) self.shape = scale.shape assert self.shape == () - self.scale = Parameterize(lambda x: x + jnp.sqrt(1 + x**2), jnp.zeros(())) - self.theta = Parameterize(lambda x: x + jnp.sqrt(1 + x**2), jnp.zeros(())) + self.scale = Parameterize(_positive, jnp.zeros(())) + self.theta = Parameterize(_positive, jnp.zeros(())) def _log_derivative_f(self, x, mu, sigma, theta): abs_x = jnp.abs(x) @@ -473,6 +484,31 @@ def inverse_and_log_det(self, y: Array, condition: Array | None = None): return self._householder(y), jnp.zeros(()) +class AutoFlow(AbstractBijection): + shape: tuple[int, ...] + params: Array + transform_and_log_det_fn: Callable + inverse_and_log_det_fn: Callable + cond_shape = None + + def __init__( + self, params: ArrayLike, shape, transform_and_log_det_fn, inverse_and_log_det_fn + ): + params = arraylike_to_array(params) + if params.ndim != 1: + raise ValueError("params must be a vector.") + self.shape = shape + self.params = params + self.transform_and_log_det_fn = transform_and_log_det_fn + self.inverse_and_log_det_fn = inverse_and_log_det_fn + + def transform_and_log_det(self, x: jnp.ndarray, condition: Array | None = None): + return self.transform_and_log_det_fn(x, self.params) + + def inverse_and_log_det(self, y: Array, condition: Array | None = None): + return self.inverse_and_log_det_fn(y, self.params) + + class MvScale(bijections.AbstractBijection): shape: tuple[int, ...] params: Array @@ -952,8 +988,8 @@ def make_single_hh(key, idx): def make_elemwise_trafo(key, n_dim, *, count=1, vmap=True): def make_elemwise(key, loc): key1, key2 = jax.random.split(key) - scale = Parameterize(lambda x: x + jnp.sqrt(1 + x**2), jnp.zeros(())) - theta = Parameterize(lambda x: x + jnp.sqrt(1 + x**2), jnp.zeros(())) + scale = Parameterize(_positive, jnp.zeros(())) + theta = Parameterize(_positive, jnp.zeros(())) affine = AsymmetricAffine( loc, @@ -1422,7 +1458,7 @@ def make_transformer( if affine_transformer: affine = bijections.Affine(jnp.zeros(()), jnp.ones(())) - scale = Parameterize(lambda x: x + jnp.sqrt(1 + x**2), jnp.zeros(())) + scale = Parameterize(_positive, jnp.zeros(())) affine = eqx.tree_at( where=lambda aff: aff.scale, pytree=affine, @@ -1432,8 +1468,8 @@ def make_transformer( if asymmetric_transformer: for loc in [0.0]: - scale = Parameterize(lambda x: x + jnp.sqrt(1 + x**2), jnp.zeros(())) - theta = Parameterize(lambda x: x + jnp.sqrt(1 + x**2), jnp.zeros(())) + scale = Parameterize(_positive, jnp.zeros(())) + theta = Parameterize(_positive, jnp.zeros(())) affine = AsymmetricAffine( jnp.zeros(()) + loc, @@ -1878,6 +1914,7 @@ def make_flow( sandwich_householder=False, activation=None, reuse_embed=False, + auto_flow=None, ): if activation is None: activation = jax.nn.leaky_relu @@ -1901,23 +1938,36 @@ def make_flow( n_draws, n_dim = positions.shape assert positions.shape == gradients.shape + # Push positions and gradients through autoflow + if auto_flow is not None: + from nutpie.transform_adapter import inverse_gradient_and_val + + positions, gradients, _logp = eqx.filter_vmap( + inverse_gradient_and_val, in_axes=(None, 0, 0, 0) + )(auto_flow, positions, gradients, jnp.zeros((n_draws,))) + positions = np.asarray(positions) + gradients = np.asarray(gradients) + if n_draws == 0: raise ValueError("No draws") elif n_draws == 1: assert np.all(gradients != 0) - diag = np.clip(1 / jnp.sqrt(jnp.abs(gradients[0])), 1e-8, 1e8) + diag = np.clip(1 / np.sqrt(np.abs(gradients[0])), 1e-8, 1e8) assert np.isfinite(diag).all() - mean = jnp.zeros_like(diag) + mean = np.zeros_like(diag) else: + # numpy, not jnp: these run once per adaptation window on a + # window-sized array, and each new draw count would otherwise cost a + # fresh XLA compile. pos_std = np.clip(positions.std(0), 1e-8, 1e8) grad_std = np.clip(gradients.std(0), 1e-8, 1e8) - diag = jnp.sqrt(pos_std / grad_std) + diag = np.sqrt(pos_std / grad_std) mean = positions.mean(0) + gradients.mean(0) * diag * diag key = jax.random.key(seed % (2**63), impl="threefry2x32") diag_param = Parameterize( - lambda x: x + jnp.sqrt(1 + x**2), + _positive, (diag**2 - 1) / (2 * diag), ) diag_affine = bijections.Affine(mean, diag) @@ -1927,9 +1977,9 @@ def make_flow( replace=diag_param, ) - flows = [ - diag_affine, - ] + flows = [diag_affine] + if auto_flow is not None: + flows.append(auto_flow) if n_layers == 0: return bijections.Chain(flows) @@ -2071,11 +2121,11 @@ def extend_flow( if True: scale = Parameterize( - lambda x: x + jnp.sqrt(1 + x**2), + _positive, jnp.array(0.0), ) theta = Parameterize( - lambda x: x + jnp.sqrt(1 + x**2), + _positive, jnp.array(0.0), ) @@ -2152,7 +2202,7 @@ def extend_flow( if False: scale = Parameterize( - lambda x: x + jnp.sqrt(1 + x**2), + _positive, jnp.array(0.0), ) affine = eqx.tree_at( @@ -2204,7 +2254,7 @@ def extend_flow( new_layer = bijections.Sandwich(inner, permute) scale = Parameterize( - lambda x: x + jnp.sqrt(1 + x**2), + _positive, jnp.zeros(n_dim), ) affine = eqx.tree_at( @@ -2227,7 +2277,7 @@ def extend_flow( ), ) scale = Parameterize( - lambda x: x + jnp.sqrt(1 + x**2), + _positive, jnp.zeros(n_dim), ) affine = eqx.tree_at( diff --git a/python/nutpie/transform_adapter.py b/python/nutpie/transform_adapter.py index 41d18f4..47e02d5 100644 --- a/python/nutpie/transform_adapter.py +++ b/python/nutpie/transform_adapter.py @@ -26,6 +26,7 @@ train_val_split, ) import optax +import optax.tree_utils as otu from paramax import unwrap, NonTrainable from nutpie.normalizing_flow import Coupling, Householder, Scan, extend_flow, make_flow @@ -301,6 +302,21 @@ def pull_transformer_grad(transformer_params): return (x, x_grad, logp + fwd_log_det) +def _bucket_draws(*arrays, min_size=8): + """Truncate a window (oldest draws first) to a power-of-two length. + + Every distinct draw count costs a fresh XLA compile of the loss, its + gradient and the vmapped transform. The windows grow by a handful of + draws each time, so without this each one is a new shape and adaptation + pays compile time it never amortizes. + """ + n = len(arrays[0]) + if n < min_size: + return arrays + size = 1 << (n.bit_length() - 1) + return tuple(a[-size:] for a in arrays) + + class FisherLoss: def __init__(self, gamma=None, log_inside_batch=False): self._gamma = gamma @@ -391,6 +407,81 @@ def fit_flow(key, bijection, loss_fn, draws, grads, logps, **kwargs): return fit.bijection, losses, opt_state +@eqx.filter_jit +def _run_lbfgs(params, static, loss_fn, draws, grads, logps, max_iter, gtol): + """L-BFGS to convergence, entirely inside one jit. + + The loop is a ``lax.while_loop`` rather than a python loop around a + jitted step so the whole fit stays on the device -- no host sync per + iteration, and it runs on a GPU like the rest of the flow machinery. + """ + + def loss(params): + return loss_fn(params, static, draws, grads, logps) + + optimizer = optax.lbfgs() + # Reuses the value/grad the zoom line search already computed. + value_and_grad = optax.value_and_grad_from_state(loss) + + def lbfgs_step(carry): + params, state = carry + value, grad = value_and_grad(params, state=state) + updates, state = optimizer.update( + grad, state, params, value=value, grad=grad, value_fn=loss + ) + return optax.apply_updates(params, updates), state + + def not_converged(carry): + _, state = carry + count = otu.tree_get(state, "count") + grad = otu.tree_get(state, "grad") + return (count == 0) | ((count < max_iter) & (otu.tree_l2_norm(grad) >= gtol)) + + init = (params, optimizer.init(params)) + params, state = jax.lax.while_loop(not_converged, lbfgs_step, init) + return params, otu.tree_get(state, "value") + + +def fit_flow_lbfgs( + key, + bijection, + loss_fn, + draws, + grads, + logps, + *, + max_iter=200, + gtol=1e-5, + **_ignored, +): + """Full-batch quasi-Newton fit of a low-dimensional flow. + + The auto-reparam flow has only a handful of trainable knobs (the VIP + ``h``), a smooth deterministic loss and a window that fits in one + batch, so minibatch SGD spends thousands of tiny steps where L-BFGS + needs a few dozen full-batch evaluations. Not for the neural flows: + those have too many parameters and rely on the stochasticity. + """ + flow = flowjax.flows.Transformed( + flowjax.distributions.StandardNormal(bijection.shape), bijection + ) + params, static = eqx.partition( + flow, + eqx.is_inexact_array, + is_leaf=lambda leaf: isinstance(leaf, NonTrainable), + ) + if not any(p.size for p in jax.tree.leaves(params)): + return bijection, {"train": [], "val": []}, None + + draws, grads, logps = (jnp.asarray(a) for a in (draws, grads, logps)) + params, value = _run_lbfgs( + params, static, loss_fn, draws, grads, logps, max_iter, gtol + ) + fit = eqx.combine(params, static) + losses = [float(value)] + return fit.bijection, {"train": losses, "val": losses}, None + + @eqx.filter_jit def _init_from_transformed_position(logp_fn, bijection, transformed_position): bijection = unwrap(bijection) @@ -463,6 +554,8 @@ def _inv_transform(bijection, untransformed_position, untransformed_gradient): class TransformAdapter: + """Does optimization""" + def __init__( self, seed, @@ -493,7 +586,11 @@ def __init__( make_optimizer=None, num_layers=9, max_epochs=200, + fit_method="sgd", ): + if fit_method not in ("sgd", "lbfgs"): + raise ValueError(f"Unknown fit_method: {fit_method}") + self._fit_method = fit_method self._logp_fn = logp_fn self._make_flow_fn = make_flow_fn self._chain = chain @@ -569,9 +666,9 @@ def update(self, seed, positions, gradients, logps): gradients = gradients_slice logps = logp_slice - positions = np.array(positions) - gradients = np.array(gradients) - logps = np.array(logps) + positions, gradients, logps = _bucket_draws( + np.array(positions), np.array(gradients), np.array(logps) + ) fit = self._make_flow_fn(seed, positions, gradients, n_layers=0) @@ -590,9 +687,11 @@ def update(self, seed, positions, gradients, logps): return - positions = np.array(positions[self._initial_skip :][-self._window_size :]) - gradients = np.array(gradients[self._initial_skip :][-self._window_size :]) - logps = np.array(logps[self._initial_skip :][-self._window_size :]) + positions, gradients, logps = _bucket_draws( + np.array(positions[self._initial_skip :][-self._window_size :]), + np.array(gradients[self._initial_skip :][-self._window_size :]), + np.array(logps[self._initial_skip :][-self._window_size :]), + ) if len(positions) < 10: return @@ -609,6 +708,7 @@ def update(self, seed, positions, gradients, logps): # TODO don't reuse seed key = jax.random.PRNGKey(seed % (2**63)) + diag_was_refit = False if len(self._bijection.bijections) == 1: base = self._make_flow_fn( seed, @@ -633,6 +733,36 @@ def update(self, seed, positions, gradients, logps): logps[-128:], ), ) + elif isinstance(self._bijection.bijections[0], bijections.Affine): + # Diag-leading chain (the auto_flow/VIP case): no trained + # component consumes the diag's output, so the closed-form + # per-dimension scale estimate stays available all through + # tuning. Re-estimate the diag from this window's draws, + # conditional on the current trailing flows (make_flow pushes + # draws/grads through ``auto_flow`` before fitting), freeze + # it, and let the optimizer train only the trailing (VIP) + # parameters. Without this, the diag would stay stuck at its + # last diag-window fit and only creep by SGD. + # std <- auto_reparam <- diag_affine <- draws + # std <- diag_affine <- nf_flows <- draws + # (code order: VIP std <- diag <- auto_reparam <- draws; + # neural std <- nf_flows <- diag <- draws) + rest = list(self._bijection.bijections[1:]) + tail = rest[0] if len(rest) == 1 else bijections.Chain(rest) + fresh = self._make_flow_fn( + seed, positions, gradients, n_layers=0, auto_flow=tail + ) + diag = fresh.bijections[0] + diag = eqx.tree_at( + lambda d: (d.loc, d.scale), + diag, + replace=(NonTrainable(diag.loc), NonTrainable(diag.scale)), + ) + base = bijections.Chain([diag, *rest]) + # Param structure changed (diag frozen out) — never reuse + # optimizer state across it. + self._opt_state = None + diag_was_refit = True else: base = self._bijection @@ -670,8 +800,8 @@ def update(self, seed, positions, gradients, logps): return flow = flowjax.flows.Transformed( - flowjax.distributions.StandardNormal(self._bijection.shape), - self._bijection, + flowjax.distributions.StandardNormal(base.shape), + base, ) params, static = eqx.partition(flow, eqx.is_inexact_array) @@ -679,12 +809,27 @@ def update(self, seed, positions, gradients, logps): params, static, positions[-128:], gradients[-128:], logps[-128:] ) - if np.isfinite(old_loss) and old_loss < -4 and self.index > 10: + # The refit diag is an unconditional improvement candidate: adopt + # it now so it survives even if the SGD step below is rejected. + if base is not self._bijection and np.isfinite(old_loss): + self._bijection = base + + # The absolute low-loss skip is miscalibrated on refit windows: + # a fresh diag alone often reaches loss < -4 while the VIP knobs + # still have large untapped gains (h stuck near 0), so only skip + # when the transform carried over unchanged. + if ( + np.isfinite(old_loss) + and old_loss < -4 + and self.index > 10 + and not diag_was_refit + ): if self._verbose: print(f"Loss is low ({old_loss}), skipping training") return - fit, _, opt_state = fit_flow( + fit_fn = fit_flow_lbfgs if self._fit_method == "lbfgs" else fit_flow + fit, _, opt_state = fit_fn( key, base, self._loss_fn, @@ -874,7 +1019,7 @@ def make_transform_adapter( show_progress=False, nn_depth=None, nn_width=None, - num_layers=8, + num_layers=None, num_diag_windows=6, learning_rate=5e-4, untransformed_dim=None, @@ -905,10 +1050,28 @@ def make_transform_adapter( contract_transformer=True, asymmetric_transformer=False, reuse_embed=True, + auto_flow=None, + fit_method=None, ): if extension_windows is None: extension_windows = [] + # Several auto flows compose into a single bijection; flowjax applies + # Chain members in order in the transform (sampler -> value) direction. + if isinstance(auto_flow, (list, tuple)): + auto_flow = bijections.Chain(list(auto_flow)) if auto_flow else None + + # With an auto flow, default to fitting only the reparametrization (and + # the diag affine); set num_layers explicitly to add coupling layers. + if num_layers is None: + num_layers = 0 if auto_flow is not None else 8 + + # A pure auto flow has only the few VIP knobs to fit: full-batch L-BFGS + # converges in a few dozen loss evaluations where SGD needs thousands of + # steps. Neural layers keep the SGD path. + if fit_method is None: + fit_method = "lbfgs" if (auto_flow is not None and num_layers == 0) else "sgd" + return partial( TransformAdapter, verbose=verbose, @@ -930,6 +1093,7 @@ def make_transform_adapter( contract_transformer=contract_transformer, asymmetric_transformer=asymmetric_transformer, reuse_embed=reuse_embed, + auto_flow=auto_flow, ), show_progress=show_progress, num_diag_windows=num_diag_windows, @@ -950,4 +1114,5 @@ def make_transform_adapter( make_optimizer=make_optimizer, num_layers=num_layers, max_epochs=max_epochs, + fit_method=fit_method, ) diff --git a/python/nutpie/transform_adapter_numba.py b/python/nutpie/transform_adapter_numba.py new file mode 100644 index 0000000..a39c692 --- /dev/null +++ b/python/nutpie/transform_adapter_numba.py @@ -0,0 +1,335 @@ +"""Pytensor/numba transform adapter for the automatic reparametrization. + +:mod:`nutpie.transform_adapter` fits general normalizing flows and needs +jax/flowjax for it. The auto-reparam flow is different: its transform is +already a pytensor graph (see :mod:`nutpie.flow_reparam`) and it has only a +handful of trainable knobs (the VIP ``h``). So the whole adapter compiles to +numba and the fit is a full-batch quasi-Newton solve: + +* The per-leapfrog-step hook (``init_from_transformed_position``, which nuts-rs + calls for *every* gradient evaluation) is one numba function computing + ``logp(constrain(z)) + logdet`` and its gradient — tens of microseconds, + instead of a python -> jax round trip per step. +* Fitting ``h`` is scipy's L-BFGS-B on the full window, a few dozen loss + evaluations, instead of thousands of minibatch SGD steps. + +The chain, sampler space to value space, is ``z --diag--> y --flow--> x``: +a diagonal affine (estimated in closed form from each window, as in the jax +adapter) followed by the reparametrization flow. +""" + +from __future__ import annotations + +import warnings +from functools import partial + +import numpy as np +import pytensor.tensor as pt +import scipy.optimize +from pytensor.gradient import pullback +from pytensor.graph import rewrite_graph +from pytensor.graph.replace import graph_replace, vectorize_graph + +from nutpie.flow_reparam import NoFlow, build_flow_graph_from_specs, free_vars_info +from nutpie.flow_reparam import reparametrize + + +def _model_logp_graph(model, infos, n_dim): + """The model's logp as a graph over nutpie's flat value vector.""" + logp = rewrite_graph(model.logp(), include=["canonicalize", "stabilize"]) + joined = pt.TensorType("float64", shape=(n_dim,))("_value_point") + by_name = {info.name: info for info in infos} + replacements = {} + for rv in model.free_RVs: + value = model.rvs_to_values[rv] + info = by_name[value.name] + chunk = joined[info.start_idx : info.end_idx] + replacements[value] = chunk.reshape(tuple(info.shape)).astype(rv.dtype) + (logp,) = graph_replace([logp], replacements) + return joined, logp + + +class NumbaAutoFlow: + """The compiled numba functions the adapter drives. + + Built once per model; the adapter instances (one per chain) hold only the + fitted parameters (``h``, and the diagonal affine's ``loc``/``scale``). + """ + + def __init__(self, model, compiled_model, specs): + from pymc.pytensorf import compile as compile_pymc + + n_dim = int(compiled_model.n_dim) + infos = free_vars_info(compiled_model) + graphs = build_flow_graph_from_specs(specs, infos, n_dim) + (y_c, h_c), (x_of_y, ljd_c) = graphs["constrain"] + (x_u, h_u), (y_of_x, ljd_u) = graphs["unconstrain"] + + self.n_dim = n_dim + self.n_params = int(graphs["flow_params_vector"].type.shape[0]) + + h = pt.dvector("h") + loc = pt.dvector("loc") + scale = pt.dvector("scale") + + # ---- sampler space -> value space -------------------------------- + z = pt.dvector("z") + x_of_z, ljd_flow = graph_replace( + [x_of_y, ljd_c], {y_c: loc + scale * z, h_c: h} + ) + ljd_of_z = ljd_flow + pt.log(scale).sum() + + joined, logp_of_joined = _model_logp_graph(model, infos, n_dim) + grad_of_joined = pt.grad(logp_of_joined, joined) + + logp_z, gx_z = graph_replace([logp_of_joined, grad_of_joined], {joined: x_of_z}) + # The transformed gradient is just d/dz of the transformed logp. + gz = pt.grad(logp_z + ljd_of_z, z) + + # ---- value space -> sampler space -------------------------------- + x = pt.dvector("x") + gx = pt.dvector("gx") + y_of_x_h, _ = graph_replace([y_of_x, ljd_u], {x_u: x, h_u: h}) + z_of_x = (y_of_x_h - loc) / scale + # logdet and transformed gradient are those of the *forward* map at the + # point we just landed on, so the sampler's bookkeeping stays in one + # convention (matching the jax adapter's inverse_gradient_and_val). + x_rec, ljd_rec = graph_replace([x_of_z, ljd_of_z], {z: z_of_x}) + gz_of_x = pullback([x_rec, ljd_rec], z_of_x, [gx, pt.ones(())]) + + logp_x, gx_x = graph_replace([logp_of_joined, grad_of_joined], {joined: x}) + gz_of_x_own = pullback([x_rec, ljd_rec], z_of_x, [gx_x, pt.ones(())]) + + # ---- the Fisher divergence over a window of draws ----------------- + X = pt.dmatrix("X") + GX = pt.dmatrix("GX") + Y_b = vectorize_graph(y_of_x, {x_u: X, h_u: h}) + Z_b = (Y_b - loc) / scale + X_rec_b = vectorize_graph(x_of_z, {z: Z_b, h: h, loc: loc, scale: scale}) + LJD_b = vectorize_graph(ljd_of_z, {z: Z_b, h: h, loc: loc, scale: scale}) + GZ_b = pullback([X_rec_b, LJD_b], Z_b, [GX, pt.ones(X.shape[0])]) + loss = pt.log(pt.sum((Z_b + GZ_b) ** 2, axis=1).mean()) + + # ---- draws pushed into flow space, for the diagonal estimate ------ + X_rec_y = vectorize_graph([x_of_y, ljd_c], {y_c: Y_b, h_c: h}) + GY_b = pullback(list(X_rec_y), Y_b, [GX, pt.ones(X.shape[0])]) + + def build(inputs, outputs): + # The jitted function directly, not the pytensor Function wrapper: + # the wrapper's input validation and storage bookkeeping is a large + # share of the cost of a call this small, and the hot hook runs once + # per leapfrog step. Shared variables (e.g. pm.Data) become trailing + # arguments of the jitted signature. + fn = compile_pymc(inputs, outputs, mode="NUMBA") + jit_fn = fn.vm.jit_fn + shared = [var.get_value(borrow=True) for var in fn.get_shared()] + if not shared: + return jit_fn + + def call(*args): + return jit_fn(*args, *shared) + + return call + + # Called once per leapfrog step -- everything else is per window. + self.transformed = build( + [z, h, loc, scale], [logp_z, ljd_of_z, x_of_z, gx_z, gz] + ) + self.inv_transform = build([x, gx, h, loc, scale], [ljd_rec, z_of_x, gz_of_x]) + self.untransformed = build( + [x, h, loc, scale], [logp_x, ljd_rec, gx_x, z_of_x, gz_of_x_own] + ) + self.loss_and_grad = build([X, GX, h, loc, scale], [loss, pt.grad(loss, h)]) + self.push_to_flow_space = build([X, GX, h], [Y_b, GY_b]) + + +class NumbaTransformAdapter: + """One per chain; nuts-rs calls these methods (see ``src/wrapper.rs``).""" + + def __init__( + self, + seed, + position, + gradient, + chain, + *, + flow: NumbaAutoFlow, + window_size=600, + num_diag_windows=6, + initial_skip=120, + max_iter=200, + verbose=False, + ): + self._flow = flow + self._chain = chain + self._window_size = window_size + self._num_diag_windows = num_diag_windows + self._initial_skip = initial_skip + self._max_iter = max_iter + self._verbose = verbose + + self._h = np.zeros(flow.n_params) + self._loc = np.zeros(flow.n_dim) + self._scale = np.ones(flow.n_dim) + self.index = 0 + + # h = 0 is the centred no-op, so the single draw needs no push-through. + gradient = np.asarray(gradient, dtype="float64") + with np.errstate(divide="ignore"): + scale = 1 / np.sqrt(np.abs(gradient)) + if np.isfinite(scale).all(): + self._scale = np.clip(scale, 1e-8, 1e8) + + @property + def transformation_id(self): + return self.index + + def _fit_diag(self, positions, gradients): + """Closed-form per-dimension affine, conditional on the current flow.""" + y, gy = self._flow.push_to_flow_space(positions, gradients, self._h) + pos_std = np.clip(y.std(0), 1e-8, 1e8) + grad_std = np.clip(gy.std(0), 1e-8, 1e8) + scale = np.sqrt(pos_std / grad_std) + loc = y.mean(0) + gy.mean(0) * scale * scale + return loc, scale + + def _loss(self, positions, gradients, h, loc, scale): + value, _ = self._flow.loss_and_grad(positions, gradients, h, loc, scale) + return float(value) + + def update(self, seed, positions, gradients, logps): + self.index += 1 + if len(positions) == 0: + return + + positions = np.ascontiguousarray(positions, dtype="float64") + gradients = np.ascontiguousarray(gradients, dtype="float64") + + # Early windows: only re-estimate the diagonal, as the jax adapter does. + if self.index <= self._num_diag_windows: + size = len(positions) + lower = -size // 5 + 3 + if len(positions[lower:]) > 0: + positions, gradients = positions[lower:], gradients[lower:] + loc, scale = self._fit_diag(positions, gradients) + if np.isfinite(loc).all() and np.isfinite(scale).all(): + self._loc, self._scale = loc, scale + return + + # Numba dispatches on types, not shapes, so a window of any length runs + # without recompiling -- no need to pad or truncate to fixed sizes. + positions = positions[self._initial_skip :][-self._window_size :] + gradients = gradients[self._initial_skip :][-self._window_size :] + if len(positions) < 10: + return + if not (np.isfinite(positions).all() and np.isfinite(gradients).all()): + return + + # Re-estimate the diagonal from this window (conditional on the current + # flow), then fit the flow knobs with the diagonal held fixed. + loc, scale = self._fit_diag(positions, gradients) + if not (np.isfinite(loc).all() and np.isfinite(scale).all()): + loc, scale = self._loc, self._scale + + old_loss = self._loss(positions, gradients, self._h, loc, scale) + if np.isfinite(old_loss): + self._loc, self._scale = loc, scale + + def fun(h): + value, grad = self._flow.loss_and_grad( + positions, gradients, np.ascontiguousarray(h), loc, scale + ) + value = float(value) + if not np.isfinite(value): + # Abort the line search, not the fit: scipy backtracks. + return np.inf, np.zeros_like(h) + return value, np.asarray(grad, dtype="float64") + + result = scipy.optimize.minimize( + fun, + self._h, + jac=True, + method="L-BFGS-B", + options={"maxiter": self._max_iter}, + ) + new_loss = self._loss(positions, gradients, result.x, loc, scale) + + if self._verbose: + print( + f"Chain {self._chain} window {self.index}: " + f"loss {old_loss:.4f} -> {new_loss:.4f} in {result.nfev} evals" + ) + + if np.isfinite(new_loss) and (not np.isfinite(old_loss) or new_loss < old_loss): + self._h = np.ascontiguousarray(result.x) + + def init_from_transformed_position(self, transformed_position): + logp, logdet, x, gx, gz = self._flow.transformed( + np.ascontiguousarray(transformed_position, dtype="float64"), + self._h, + self._loc, + self._scale, + ) + return float(logp), float(logdet), x, gx, gz + + def init_from_untransformed_position(self, untransformed_position): + logp, logdet, gx, z, gz = self._flow.untransformed( + np.ascontiguousarray(untransformed_position, dtype="float64"), + self._h, + self._loc, + self._scale, + ) + return float(logp), float(logdet), gx, z, gz + + def inv_transform(self, position, gradient): + logdet, z, gz = self._flow.inv_transform( + np.ascontiguousarray(position, dtype="float64"), + np.ascontiguousarray(gradient, dtype="float64"), + self._h, + self._loc, + self._scale, + ) + return float(logdet), z, gz + + +def make_numba_transform_adapter( + *, + numba_flow: NumbaAutoFlow, + window_size=600, + num_diag_windows=6, + initial_skip=120, + max_iter=200, + verbose=False, + **_ignored, +): + return partial( + NumbaTransformAdapter, + flow=numba_flow, + window_size=window_size, + num_diag_windows=num_diag_windows, + initial_skip=initial_skip, + max_iter=max_iter, + verbose=verbose, + ) + + +def build_auto_flow_numba(model, compiled_model): + """Compile ``model``'s auto-reparam flow for the numba adapter. + + Returns ``None`` (with a warning) if the rewrite found nothing to + reparametrize, mirroring :func:`nutpie.flow_reparam.build_auto_flow`. + """ + specs = reparametrize(model) + flowed = [s for s in specs if s.flow_cls is not NoFlow] + if not flowed: + warnings.warn( + "Automatic reparametrization did not find any variables to " + "reparametrize in this model." + ) + return None + chosen = ", ".join(f"{s.value.name} ({s.flow_cls.__name__})" for s in flowed) + print( + f"auto_reparam: reparametrizing {len(flowed)} of {len(specs)} " + f"free variables: {chosen}" + ) + return NumbaAutoFlow(model, compiled_model, specs) diff --git a/tests/test_flow_reparam.py b/tests/test_flow_reparam.py new file mode 100644 index 0000000..f31af41 --- /dev/null +++ b/tests/test_flow_reparam.py @@ -0,0 +1,970 @@ +from importlib.util import find_spec + +import pytest + +if find_spec("pymc") is None: + pytest.skip("Skip pymc tests", allow_module_level=True) + +import numpy as np +import pymc as pm +import pytensor +import pytensor.tensor as pt +from pymc import dims as pmd +from pymc.distributions.transforms import Interval + +try: + from pytensor.gradient import pullback +except ImportError: # pytensor < 3.0 used the name Lop + from pytensor.gradient import Lop as pullback + +import nutpie +from nutpie.flow_reparam import ( + AffineFlow, + NoFlow, + ShiftFlow, + automatic_flow_reparam, + build_auto_flow, + build_flow_graph, + free_vars_info, +) + + +def _compile_flow(model): + """Compile the model with nutpie, build the flow graph over its + flat variables, and compile the pytensor constrain/unconstrain + functions.""" + records = automatic_flow_reparam(model) + compiled = nutpie.compile_pymc_model(model, backend="jax", gradient_backend="jax") + g = build_flow_graph(model, free_vars_info(compiled), compiled.n_dim) + constrain_fn = pytensor.function(*g["constrain"]) + unconstrain_fn = pytensor.function(*g["unconstrain"]) + return compiled, records, constrain_fn, unconstrain_fn + + +def _total_params(records): + return sum(int(np.prod(sh)) for r in records.values() for sh in r["param_shapes"]) + + +def _var_slice(compiled, name): + v = next(v for v in compiled._variables if v.name == name) + return slice(v.start_idx, v.end_idx) + + +@pytest.mark.pymc +def test_root_rv_not_reparametrized(): + with pm.Model() as m: + pm.Normal("x", 0, 1, shape=(3,)) + + records = automatic_flow_reparam(m) + assert records["x"]["flow_cls"] is NoFlow + assert records["x"]["param_shapes"] == [] + + +@pytest.mark.pymc +def test_transformed_rv_not_reparametrized(): + coords = {"group": [0, 1, 2]} + with pm.Model(coords=coords) as m: + pop_mu = pm.Normal("pop_mu", 0, 1) + pop_sigma = pm.HalfNormal("pop_sigma", 1) + # This could be fine, but there may rewrites + # like ordered/zerosum that change things too much? + pm.Normal( + "ind_mu", + pop_mu, + pop_sigma, + dims="group", + transform=Interval(lower=-10, upper=10), + ) + + records = automatic_flow_reparam(m) + (name,) = [n for n in records if n.startswith("ind_mu")] + assert records[name]["flow_cls"] is NoFlow + + +@pytest.mark.parametrize( + "dist_fn", + [ + lambda mu, sigma, **k: pm.Normal("y", mu, sigma, **k), + lambda mu, sigma, **k: pm.Cauchy("y", mu, sigma, **k), + lambda mu, sigma, **k: pm.Laplace("y", mu, sigma, **k), + lambda mu, sigma, **k: pm.Logistic("y", mu, sigma, **k), + lambda mu, sigma, **k: pm.Gumbel("y", mu, sigma, **k), + lambda mu, sigma, **k: pm.StudentT("y", nu=3, mu=mu, sigma=sigma, **k), + lambda mu, sigma, **k: pm.LogNormal("y", mu, sigma, **k), + ], + ids=["Normal", "Cauchy", "Laplace", "Logistic", "Gumbel", "StudentT", "LogNormal"], +) +@pytest.mark.pymc +def test_loc_scale_affine_flow(dist_fn): + with pm.Model(coords={"group": range(4)}) as m: + mu = pm.Normal("mu", 0, 1) + sigma = pm.HalfNormal("sigma", 1) + dist_fn(mu=mu, sigma=sigma, dims="group") + + records = automatic_flow_reparam(m) + (name,) = [n for n in records if n in {"y", "y_log__"}] + assert records[name]["flow_cls"] is AffineFlow + assert records[name]["param_shapes"] == [(4,), (4,)] + + +@pytest.mark.pymc +def test_per_param_qualification(): + """A dist param earns a per-element knob iff it depends on another + latent; a param that is constant (or data-only) gets a size-0 knob, + pinned at the centred no-op — regardless of how it maps onto the RV's + axes (broadcast hyperprior or 1:1 full-shape).""" + coords = {"group": range(3)} + + with pm.Model(coords=coords) as m: + sigma = pm.HalfNormal("sigma", 1) + pm.Normal("y", 0.0, sigma, dims="group") + r = automatic_flow_reparam(m)["y"] + assert r["flow_cls"] is AffineFlow + assert r["param_shapes"] == [(0,), (3,)] + + with pm.Model(coords=coords) as m: + mu = pm.Normal("mu", 0, 1) + pm.Normal("y", mu, 2.0, dims="group") + r = automatic_flow_reparam(m)["y"] + assert r["flow_cls"] is AffineFlow + assert r["param_shapes"] == [(3,), (0,)] + + # Full-shape (1:1 elementwise) scale that still depends on a latent — + # each y_i funnels off its own sigma_i — earns a per-element knob. + with pm.Model(coords=coords) as m: + mu = pm.Normal("mu", 0, 1) + sigma = pm.HalfNormal("sigma", 1, dims="group") + pm.Normal("y", mu, sigma, dims="group") + r = automatic_flow_reparam(m)["y"] + assert r["flow_cls"] is AffineFlow + assert r["param_shapes"] == [(3,), (3,)] + + with pm.Model(coords=coords) as m: + pm.Normal("root", 0, 1) + pm.Normal("y", 0.0, 2.0, dims="group") + r = automatic_flow_reparam(m)["y"] + assert r["flow_cls"] is NoFlow + + +@pytest.mark.pymc +def test_scalar_funnel_reparametrized(): + """Neal's funnel: a scalar latent whose scale funnels off another + scalar latent fires, with the knob itself a scalar (shape ()). The + dependence is 1:1 with no broadcast axis, and the value is 0-d — both + used to be skipped.""" + with pm.Model() as m: + z = pm.Normal("z", 0.0, 3.0) + pm.Normal("x", 0.0, pt.exp(z / 2)) + + compiled, records, constrain_fn, unconstrain_fn = _compile_flow(m) + assert records["z"]["flow_cls"] is NoFlow + assert records["x"]["flow_cls"] is AffineFlow + # loc = 0 withheld; the scale knob is a scalar. + assert records["x"]["param_shapes"] == [(0,), ()] + + n_dim = int(compiled.n_dim) + assert n_dim == 2 + z_sl = _var_slice(compiled, "z") + x_sl = _var_slice(compiled, "x") + + rng = np.random.default_rng(0) + y = rng.normal(size=n_dim) + h_sigma = 0.6 + flow_params = np.array([h_sigma]) + + value, ljd_c = constrain_fn(y, flow_params) + sigma = np.exp(y[z_sl][0] / 2) + np.testing.assert_allclose(value[x_sl], sigma**h_sigma * y[x_sl], atol=1e-10) + np.testing.assert_allclose(ljd_c, h_sigma * np.log(sigma), atol=1e-10) + + y_back, ljd_u = unconstrain_fn(value, flow_params) + np.testing.assert_allclose(y_back, y, atol=1e-10) + np.testing.assert_allclose(ljd_c + ljd_u, 0.0, atol=1e-10) + + +@pytest.mark.pymc +def test_scale_shift_flow(): + coords = {"group": [0, 1, 2]} + with pm.Model(coords=coords) as m: + beta = pm.HalfNormal("beta", 1) + pm.Gamma("x", alpha=2.0, beta=beta, dims="group") + + r = automatic_flow_reparam(m)["x_log__"] + assert r["flow_cls"] is ShiftFlow + assert r["param_shapes"] == [(3,)] + + +@pytest.mark.pymc +def test_zerosum_scale_flow(): + with pm.Model(coords={"group": range(5)}) as m: + sigma = pm.HalfNormal("pop_sigma", 1) + pm.ZeroSumNormal("x", sigma=sigma, dims="group") + + r = automatic_flow_reparam(m)["x_zerosum__"] + assert r["flow_cls"] is AffineFlow + # No loc knob; one scale knob shared across the zero-sum core dim. + assert r["param_shapes"] == [(0,), (1,)] + + with pm.Model(coords={"batch": range(2), "group": range(5)}) as m: + sigma = pm.HalfNormal("pop_sigma", 1) + pm.ZeroSumNormal("x", sigma=sigma, dims=("batch", "group")) + + r = automatic_flow_reparam(m)["x_zerosum__"] + assert r["flow_cls"] is AffineFlow + assert r["param_shapes"] == [(0,), (2, 1)] + + with pm.Model(coords={"group": range(5)}) as m: + pm.Normal("root", 0, 1) + pm.ZeroSumNormal("x", sigma=2.0, dims="group") + + r = automatic_flow_reparam(m)["x_zerosum__"] + assert r["flow_cls"] is NoFlow + + +@pytest.mark.pymc +def test_hierarchical_normal(): + coords = {"group": [0, 1, 2]} + with pm.Model(coords=coords) as m: + pop_mu = pm.Normal("pop_mu", 0, 1) + pop_sigma = pm.HalfNormal("pop_sigma", 1) + pm.Normal("ind_mu", pop_mu, pop_sigma, dims="group") + + compiled, records, constrain_fn, unconstrain_fn = _compile_flow(m) + + assert records["pop_mu"]["flow_cls"] is NoFlow + assert records["pop_sigma_log__"]["flow_cls"] is NoFlow + assert records["ind_mu"]["flow_cls"] is AffineFlow + assert records["ind_mu"]["param_shapes"] == [(3,), (3,)] + + n_dim = int(compiled.n_dim) + assert n_dim == 5 # pop_mu + pop_sigma_log__ + ind_mu(group=3) + total_params = _total_params(records) + assert total_params == 6 + pop_mu_sl = _var_slice(compiled, "pop_mu") + pop_sigma_sl = _var_slice(compiled, "pop_sigma_log__") + ind_mu_sl = _var_slice(compiled, "ind_mu") + + # Centred no-op (h = 0): both directions are the identity, log|J| = 0. + point = np.arange(n_dim, dtype="float64") + zero_params = np.zeros(total_params, dtype="float64") + c_point, ljd_c = constrain_fn(point, zero_params) + u_point, ljd_u = unconstrain_fn(point, zero_params) + np.testing.assert_allclose(c_point, point, atol=1e-9) + np.testing.assert_allclose(u_point, point, atol=1e-9) + np.testing.assert_allclose(ljd_c, 0.0, atol=1e-9) + np.testing.assert_allclose(ljd_u, 0.0, atol=1e-9) + + # Random non-identity: roundtrip and log|J| cancellation. + rng = np.random.default_rng(0) + phi0 = rng.normal(size=n_dim).astype("float64") + rand_params = rng.normal(size=total_params).astype("float64") * 0.3 + value, ljd_c = constrain_fn(phi0, rand_params) + phi_back, ljd_u = unconstrain_fn(value, rand_params) + np.testing.assert_allclose(phi_back, phi0, atol=1e-10) + np.testing.assert_allclose(ljd_c + ljd_u, 0.0, atol=1e-10) + assert not np.isclose(ljd_c, 0.0) + + # Analytical per-element VIP transform (Gorinova et al. 2019, with + # h = 1-λ so that h = 0 is centred): + # value_i = μ + σ^h_σi·(y_i - (1-h_μi)·μ), log|J| = Σ_i h_σi·log σ + h_mu = np.array([0.4, 0.2, 0.0]) + h_sigma = np.array([0.1, 0.5, 0.9]) + analytic_params = np.concatenate([h_mu, h_sigma]) + y = rng.normal(size=n_dim).astype("float64") + mu = y[pop_mu_sl][0] + sigma = np.exp(y[pop_sigma_sl][0]) + value, ljd_c = constrain_fn(y, analytic_params) + expected = mu + sigma**h_sigma * (y[ind_mu_sl] - (1 - h_mu) * mu) + np.testing.assert_allclose(value[ind_mu_sl], expected, atol=1e-10) + np.testing.assert_allclose(ljd_c, (h_sigma * np.log(sigma)).sum(), atol=1e-10) + + +@pytest.mark.pymc +def test_partial_broadcast(): + coords = {"group": range(3), "rep": range(4)} + with pm.Model(coords=coords) as m: + pop_mu = pm.Normal("pop_mu", 0, 1, dims="group") + pm.Normal("x", pop_mu[:, None], 1.0, dims=("group", "rep")) + + compiled, records, constrain_fn, unconstrain_fn = _compile_flow(m) + assert records["x"]["flow_cls"] is AffineFlow + assert records["x"]["param_shapes"] == [(3, 4), (0,)] + + n_dim = int(compiled.n_dim) + assert n_dim == 15 # pop_mu(group=3) + x(group=3, rep=4) + rng = np.random.default_rng(1) + phi0 = rng.normal(size=n_dim).astype("float64") + flow_params = rng.normal(size=_total_params(records)).astype("float64") * 0.2 + + value, ljd_c = constrain_fn(phi0, flow_params) + phi_back, ljd_u = unconstrain_fn(value, flow_params) + np.testing.assert_allclose(phi_back, phi0, atol=1e-10) + # With the scale knob pinned, the transform is a translation. + np.testing.assert_allclose(ljd_c, 0.0, atol=1e-10) + np.testing.assert_allclose(ljd_u, 0.0, atol=1e-10) + + +@pytest.mark.pymc +def test_flow_alongside_dirichlet(): + coords = {"group": [0, 1], "k": range(3)} + with pm.Model(coords=coords) as m: + pi = pm.Dirichlet("pi", a=np.ones(3), dims="k") + pop_sigma = pm.HalfNormal("pop_sigma", 1) + pm.Normal("ind_mu", pi, pop_sigma, dims=("group", "k")) + + compiled, records, constrain_fn, unconstrain_fn = _compile_flow(m) + assert records["ind_mu"]["flow_cls"] is AffineFlow + assert records["ind_mu"]["param_shapes"] == [(2, 3), (2, 3)] + (pi_name,) = [n for n in records if n.startswith("pi")] + assert records[pi_name]["flow_cls"] is NoFlow + + n_dim = int(compiled.n_dim) + assert n_dim == 9 # 2 (dirichlet, simplex-transformed) + 1 + 6 + total_params = _total_params(records) + + phi0 = np.arange(n_dim, dtype="float64") + zero_params = np.zeros(total_params, dtype="float64") + c_point, ljd_c = constrain_fn(phi0, zero_params) + np.testing.assert_allclose(c_point, phi0, atol=1e-9) + np.testing.assert_allclose(ljd_c, 0.0, atol=1e-9) + + rng = np.random.default_rng(0) + flow_params = rng.normal(size=total_params).astype("float64") * 0.3 + value, ljd_c = constrain_fn(phi0, flow_params) + phi_back, ljd_u = unconstrain_fn(value, flow_params) + np.testing.assert_allclose(phi_back, phi0, atol=1e-10) + np.testing.assert_allclose(ljd_c + ljd_u, 0.0, atol=1e-10) + + +@pytest.mark.pymc +def test_zerosum_roundtrip(): + with pm.Model(coords={"group": range(5)}) as m: + sigma = pm.HalfNormal("pop_sigma", 1) + pm.ZeroSumNormal("x", sigma=sigma, dims="group") + + compiled, records, constrain_fn, unconstrain_fn = _compile_flow(m) + n_dim = int(compiled.n_dim) + assert n_dim == 5 # pop_sigma_log__ + x_zerosum__(4) + total_params = _total_params(records) + assert total_params == 1 + + phi0 = np.arange(n_dim, dtype="float64") + zero_params = np.zeros(total_params, dtype="float64") + c_point, ljd_c = constrain_fn(phi0, zero_params) + np.testing.assert_allclose(c_point, phi0, atol=1e-9) + np.testing.assert_allclose(ljd_c, 0.0, atol=1e-9) + + rng = np.random.default_rng(0) + phi0 = rng.normal(size=n_dim).astype("float64") + flow_params = np.array([0.7]) + value, ljd_c = constrain_fn(phi0, flow_params) + phi_back, ljd_u = unconstrain_fn(value, flow_params) + np.testing.assert_allclose(phi_back, phi0, atol=1e-10) + np.testing.assert_allclose(ljd_c + ljd_u, 0.0, atol=1e-10) + + # Full NCP (h = 1): value = y·σ, log|J| = 4·log σ. + x_sl = _var_slice(compiled, "x_zerosum__") + s_sl = _var_slice(compiled, "pop_sigma_log__") + value, ljd_c = constrain_fn(phi0, np.array([1.0])) + sigma_val = np.exp(phi0[s_sl][0]) + np.testing.assert_allclose(value[x_sl], phi0[x_sl] * sigma_val, atol=1e-10) + np.testing.assert_allclose(ljd_c, 4 * np.log(sigma_val), atol=1e-10) + + +@pytest.mark.pymc +def test_dim_distributions(): + coords = {"group": [0, 1, 2]} + with pm.Model(coords=coords) as m: + pop_mu = pmd.Normal("pop_mu", 0, 1) + pop_sigma = pmd.HalfNormal("pop_sigma", 1) + pmd.Normal("ind_mu", pop_mu, pop_sigma, dims=("group",)) + pmd.LogNormal("scale", pop_mu, pop_sigma, dims=("group",)) + + compiled, records, constrain_fn, unconstrain_fn = _compile_flow(m) + assert records["pop_mu"]["flow_cls"] is NoFlow + assert records["pop_sigma_log__"]["flow_cls"] is NoFlow + assert records["ind_mu"]["flow_cls"] is AffineFlow + assert records["ind_mu"]["param_shapes"] == [(3,), (3,)] + assert records["scale_log__"]["flow_cls"] is AffineFlow + assert records["scale_log__"]["param_shapes"] == [(3,), (3,)] + + n_dim = int(compiled.n_dim) + assert n_dim == 8 # pop_mu + pop_sigma_log__ + ind_mu(3) + scale_log__(3) + total_params = _total_params(records) + phi0 = np.arange(n_dim, dtype="float64") + + zero_params = np.zeros(total_params, dtype="float64") + c_point, ljd_c = constrain_fn(phi0, zero_params) + np.testing.assert_allclose(c_point, phi0, atol=1e-9) + np.testing.assert_allclose(ljd_c, 0.0, atol=1e-9) + + rng = np.random.default_rng(0) + flow_params = rng.normal(size=total_params).astype("float64") * 0.2 + value, ljd_c = constrain_fn(phi0, flow_params) + phi_back, ljd_u = unconstrain_fn(value, flow_params) + np.testing.assert_allclose(phi_back, phi0, atol=1e-10) + np.testing.assert_allclose(ljd_c + ljd_u, 0.0, atol=1e-10) + + +def _fisher_loss_fn(model, compiled): + """Build ``loss(draws, flow_params)`` mirroring nutpie's FisherLoss: + pull value-space posterior draws and logp-gradients back through the + flow and measure deviation from a standard-normal score, minimized + analytically over the per-coordinate affine that the production + chain's diagonal-affine layers would absorb. Zero iff the pulled-back + posterior is iid normal up to a diagonal affine.""" + n_dim = int(compiled.n_dim) + free_vars = free_vars_info(compiled) + g = build_flow_graph(model, free_vars, n_dim) + + (y, flow_params), (value_out, ljd) = g["constrain"] + value_grad = value_out.type("value_grad") + # vjp of (constrain, log|J|) at cotangent (grad, 1), as in + # transform_adapter.inverse_gradient_and_val. + pulled_grad = pullback([value_out, ljd], y, [value_grad, pt.ones(())]) + pullback_fn = pytensor.function([y, flow_params, value_grad], pulled_grad) + unconstrain_fn = pytensor.function(*g["unconstrain"]) + + value_vars = {v.name: v for v in model.value_vars} + ordered = [value_vars[v.name] for v in free_vars] + grad_fn = pytensor.function(ordered, pt.grad(model.logp(), ordered)) + + def flat_grad(draw): + vals = [ + draw[v.start_idx : v.end_idx].reshape(tuple(int(s) for s in v.shape)) + for v in free_vars + ] + return np.concatenate([np.asarray(a).ravel() for a in grad_fn(*vals)]) + + def loss(draws, params): + params = np.asarray(params, dtype="float64") + xs, gs = [], [] + for draw in draws: + x, _ = unconstrain_fn(draw, params) + xs.append(x) + gs.append(pullback_fn(x, params, flat_grad(draw))) + x, g = np.array(xs), np.array(gs) + # min over per-coordinate affine z = a·(x-b) of E[(z + g/a)²]: + # 2·(sqrt(Var x·Var g) + Cov(x, g)), ≥ 0 by Cauchy-Schwarz. + cov = ((x - x.mean(0)) * (g - g.mean(0))).mean(0) + # ≥ 0 by Cauchy-Schwarz; clamp the float noise around exact zeros. + costs = np.maximum(2 * (np.sqrt(x.var(0) * g.var(0)) + cov), 0.0) + return float(np.log(np.maximum(costs.sum(), 1e-300))) + + return loss + + +def _funnel_model(obs_sigma=None, n_groups=5, n_obs=1000): + """The mixed-balance funnel from notebooks/auto_reparam-Copy1.ipynb: + per-group means with a common log-scale hyper, optionally observed + with per-group noise.""" + coords = {"group": range(n_groups)} + with pm.Model(coords=coords) as m: + s = pm.Normal("pop_sigma_log", 0, 1) + ind_mu = pm.Normal("ind_mu", 0, pm.math.exp(s / 2), dims="group") + if obs_sigma is not None: + rng = np.random.default_rng(1) + y = rng.normal( + loc=np.linspace(-0.5, 0.5, n_groups), + scale=1.0, + size=(n_obs, n_groups), + ) + pm.Normal("y", ind_mu, sigma=np.asarray(obs_sigma), observed=y) + return m, y + return m, None + + +def _funnel_posterior_draws(compiled, rng, n_draws, obs_sigma=None, y=None): + """Exact posterior draws in nutpie's flat value space: dense-grid + sampling for the 1-D hyper (conjugate marginal over the group means), + then the exact Gaussian conditional for ind_mu | s, y.""" + s_sl = _var_slice(compiled, "pop_sigma_log") + m_sl = _var_slice(compiled, "ind_mu") + n_groups = m_sl.stop - m_sl.start + + if obs_sigma is None: + s = rng.normal(0.0, 1.0, size=n_draws) + means = np.zeros((n_draws, n_groups)) + post_sd = np.exp(s / 2)[:, None] * np.ones(n_groups) + else: + obs_sigma = np.asarray(obs_sigma, dtype="float64") + n_obs = y.shape[0] + ybar = y.mean(axis=0) + grid = np.linspace(-10, 10, 8001) + tau2 = np.exp(grid) + # p(s | y) ∝ p(s) · Π_i N(ȳ_i; 0, τ² + σ_i²/n) + marg_var = tau2[:, None] + (obs_sigma**2 / n_obs)[None, :] + log_post = -0.5 * grid**2 - 0.5 * ( + np.log(marg_var) + ybar[None, :] ** 2 / marg_var + ).sum(axis=1) + p = np.exp(log_post - log_post.max()) + s = rng.choice(grid, size=n_draws, p=p / p.sum()) + prec = 1 / np.exp(s)[:, None] + (n_obs / obs_sigma**2)[None, :] + means = (n_obs * ybar / obs_sigma**2)[None, :] / prec + post_sd = 1 / np.sqrt(prec) + + ind_mu = rng.normal(means, post_sd) + draws = np.empty((n_draws, int(compiled.n_dim))) + draws[:, s_sl] = s[:, None] + draws[:, m_sl] = ind_mu + return draws + + +@pytest.mark.pymc +def test_loss_prefers_noncentered_on_prior_funnel(): + m, _ = _funnel_model() + compiled = nutpie.compile_pymc_model(m, backend="jax", gradient_backend="jax") + loss = _fisher_loss_fn(m, compiled) + + rng = np.random.default_rng(42) + draws = _funnel_posterior_draws(compiled, rng, n_draws=256) + + cp = loss(draws, np.zeros(5)) + ncp = loss(draws, np.ones(5)) + assert ncp < cp + # Full NCP standardizes the prior funnel exactly, so the loss is ~0. + assert ncp < -25 + + +@pytest.mark.pymc +def test_loss_prefers_centered_on_strong_data(): + obs_sigma = np.ones(5) + m, y = _funnel_model(obs_sigma=obs_sigma) + compiled = nutpie.compile_pymc_model(m, backend="jax", gradient_backend="jax") + loss = _fisher_loss_fn(m, compiled) + + rng = np.random.default_rng(42) + draws = _funnel_posterior_draws(compiled, rng, 256, obs_sigma=obs_sigma, y=y) + + assert loss(draws, np.zeros(5)) < loss(draws, np.ones(5)) + + +@pytest.mark.pymc +def test_loss_prefers_mixed_on_mixed_balance(): + """Strong-evidence groups want centred, weak ones non-centred: the + per-element mixed parameterization beats both global ones — the case + a single per-group knob cannot express.""" + obs_sigma = np.array([1.0, 1000.0, 1.0, 1000.0, 1.0]) + m, y = _funnel_model(obs_sigma=obs_sigma) + compiled = nutpie.compile_pymc_model(m, backend="jax", gradient_backend="jax") + loss = _fisher_loss_fn(m, compiled) + + rng = np.random.default_rng(42) + draws = _funnel_posterior_draws(compiled, rng, 256, obs_sigma=obs_sigma, y=y) + + mixed = loss(draws, np.array([0.0, 1.0, 0.0, 1.0, 0.0])) + assert mixed < loss(draws, np.zeros(5)) + assert mixed < loss(draws, np.ones(5)) + + +@pytest.mark.pymc +def test_loss_prefers_noncentered_on_zerosum_prior(): + with pm.Model(coords={"group": range(5)}) as m: + s = pm.Normal("s", 0, 1) + x = pm.ZeroSumNormal("x", sigma=pm.math.exp(s), dims="group") + + compiled = nutpie.compile_pymc_model(m, backend="jax", gradient_backend="jax") + loss = _fisher_loss_fn(m, compiled) + + s_draws, x_draws = pm.draw([s, x], draws=256, random_seed=42) + x_in = pt.matrix("x_in") + forward_fn = pytensor.function([x_in], m.rvs_to_transforms[x].forward(x_in)) + + draws = np.empty((256, int(compiled.n_dim))) + draws[:, _var_slice(compiled, "s")] = s_draws[:, None] + draws[:, _var_slice(compiled, "x_zerosum__")] = forward_fn(x_draws) + + cp = loss(draws, np.zeros(1)) + ncp = loss(draws, np.ones(1)) + assert ncp < cp + # The zero-sum transform is an isometry, so full NCP is exact. + assert ncp < -25 + + +@pytest.mark.pymc +def test_conditional_transform_parent(): + # A parent with a *conditional* transform (Interval reads the RV's + # distribution parameters) feeding a flow child: the transform must be + # applied with the actual rv inputs, not IR wrapper inputs. + coords = {"group": [0, 1, 2]} + with pm.Model(coords=coords) as m: + a = pm.Uniform("a", -1.0, 3.0) + pop_sigma = pm.HalfNormal("pop_sigma", 1) + pm.Normal("x", a, pop_sigma, dims="group") + + compiled, records, constrain_fn, unconstrain_fn = _compile_flow(m) + assert records["a_interval__"]["flow_cls"] is NoFlow + assert records["pop_sigma_log__"]["flow_cls"] is NoFlow + assert records["x"]["flow_cls"] is AffineFlow + assert records["x"]["param_shapes"] == [(3,), (3,)] + + n_dim = int(compiled.n_dim) + assert n_dim == 5 + a_sl = _var_slice(compiled, "a_interval__") + sigma_sl = _var_slice(compiled, "pop_sigma_log__") + x_sl = _var_slice(compiled, "x") + + rng = np.random.default_rng(2) + y = rng.normal(size=n_dim) + h_mu = np.array([0.7, 0.3, 0.1]) + h_sigma = np.array([0.2, 0.5, 0.8]) + flow_params = np.concatenate([h_mu, h_sigma]) + + value, ljd_c = constrain_fn(y, flow_params) + # interval backward: lower + (upper - lower) * sigmoid(value) + a_con = -1.0 + 4.0 / (1.0 + np.exp(-y[a_sl][0])) + sigma_con = np.exp(y[sigma_sl][0]) + expected_x = (y[x_sl] - (1 - h_mu) * a_con) * sigma_con**h_sigma + a_con + np.testing.assert_allclose(value[x_sl], expected_x, atol=1e-10) + np.testing.assert_allclose(ljd_c, (h_sigma * np.log(sigma_con)).sum(), atol=1e-10) + + y_back, ljd_u = unconstrain_fn(value, flow_params) + np.testing.assert_allclose(y_back, y, atol=1e-10) + np.testing.assert_allclose(ljd_c + ljd_u, 0.0, atol=1e-10) + + +@pytest.mark.pymc +def test_imputation_observed_rv_no_rng_leak(): + # Auto-imputation (observed array with missing entries) introduces observed + # RVs that carry an RNG. They must resolve to their observed data in the + # flow graph (as PyMC's logp does), leaving no RandomVariable/RNG behind — + # otherwise build_auto_flow's jit_fn gains spurious + # random_generator_shared_variable inputs and sampling fails. + from pytensor.graph.basic import ancestors + from pytensor.tensor.random.op import RandomVariable + + rng = np.random.default_rng(0) + data = rng.normal(size=12) + data[::3] = np.nan # missing entries -> PyMC auto-imputation + with pm.Model() as m: + mu = pm.Normal("mu", 0, 1) + sigma = pm.HalfNormal("sigma", 1) + pm.Normal("y", mu, sigma, observed=np.ma.masked_invalid(data)) + + records = automatic_flow_reparam(m) + (imputed,) = [n for n in records if "unobserved" in n] + assert records[imputed]["flow_cls"] is AffineFlow + + compiled = nutpie.compile_pymc_model(m, backend="jax", gradient_backend="jax") + g = build_flow_graph(m, free_vars_info(compiled), compiled.n_dim) + leftover = [ + a + for a in ancestors(g["constrain"][1]) + if a.owner is not None and isinstance(a.owner.op, RandomVariable) + ] + assert leftover == [], f"RNG-carrying RVs leaked into flow graph: {leftover}" + + # The jit_fn path (build_auto_flow) must compile without spurious RNG args. + assert build_auto_flow(m, compiled) is not None + + constrain_fn = pytensor.function(*g["constrain"]) + unconstrain_fn = pytensor.function(*g["unconstrain"]) + n_dim = int(compiled.n_dim) + npar = int(g["flow_params_vector"].type.shape[0]) + y0 = rng.normal(size=n_dim) + params = rng.normal(size=npar) * 0.3 + value, _ = constrain_fn(y0, params) + back, _ = unconstrain_fn(value, params) + np.testing.assert_allclose(back, y0, atol=1e-9) + + +@pytest.mark.pymc +def test_flow_parent_read_through_expression(): + # A *flow* parent read by its child through an expression (indexing): + # the constrain and unconstrain graphs must each compose their own + # direction of the parent flow (regression for in-place mutation of + # the shared param subgraphs between the two builds). + coords = {"group": range(3), "rep": range(4)} + with pm.Model(coords=coords) as m: + pop_mu = pm.Normal("pop_mu", 0, 1) + pop_sigma = pm.HalfNormal("pop_sigma", 1) + mu_g = pm.Normal("mu_g", pop_mu, pop_sigma, dims="group") + pm.Normal("x", mu_g[:, None], 1.0, dims=("group", "rep")) + + compiled, records, constrain_fn, unconstrain_fn = _compile_flow(m) + assert records["mu_g"]["flow_cls"] is AffineFlow + assert records["mu_g"]["param_shapes"] == [(3,), (3,)] + assert records["x"]["flow_cls"] is AffineFlow + # Constant sigma: the scale knob is withheld (size 0) and pinned. + assert records["x"]["param_shapes"] == [(3, 4), (0,)] + + n_dim = int(compiled.n_dim) + assert n_dim == 17 # 1 + 1 + 3 + 12 + total_params = _total_params(records) + assert total_params == 18 + + rng = np.random.default_rng(3) + y = rng.normal(size=n_dim) + h1_mu = rng.normal(size=3) * 0.4 + h1_sigma = rng.normal(size=3) * 0.4 + h2_mu = rng.normal(size=(3, 4)) * 0.4 + flow_params = np.concatenate([h1_mu, h1_sigma, h2_mu.ravel()]) + + mu_sl = _var_slice(compiled, "pop_mu") + sigma_sl = _var_slice(compiled, "pop_sigma_log__") + mug_sl = _var_slice(compiled, "mu_g") + x_sl = _var_slice(compiled, "x") + + value, ljd_c = constrain_fn(y, flow_params) + pop_mu_con = y[mu_sl][0] + sigma_con = np.exp(y[sigma_sl][0]) + mug_con = (y[mug_sl] - (1 - h1_mu) * pop_mu_con) * sigma_con**h1_sigma + pop_mu_con + # With the scale knob pinned the child flow is the translation + # y + h_mu·loc. + expected_x = y[x_sl].reshape(3, 4) + h2_mu * mug_con[:, None] + np.testing.assert_allclose(value[mug_sl], mug_con, atol=1e-10) + np.testing.assert_allclose(value[x_sl], expected_x.ravel(), atol=1e-10) + np.testing.assert_allclose(ljd_c, (h1_sigma * np.log(sigma_con)).sum(), atol=1e-10) + + y_back, ljd_u = unconstrain_fn(value, flow_params) + np.testing.assert_allclose(y_back, y, atol=1e-10) + np.testing.assert_allclose(ljd_c + ljd_u, 0.0, atol=1e-10) + + +@pytest.mark.pymc +def test_xtensor_unknown_transform_parent(): + # A dims RV whose transform has no plain counterpart (Beta -> logodds) + # is not lifted: it stays xtensor-typed end to end and its dim + # transform is applied natively. + coords = {"group": [0, 1, 2]} + with pm.Model(coords=coords) as m: + p = pmd.Beta("p", 1.0, 1.0) + pop_sigma = pmd.HalfNormal("pop_sigma", 1) + pmd.Normal("x", p, pop_sigma, dims=("group",)) + + compiled, records, constrain_fn, unconstrain_fn = _compile_flow(m) + (p_name,) = [n for n in records if n.startswith("p_")] + assert records[p_name]["flow_cls"] is NoFlow + assert records["x"]["flow_cls"] is AffineFlow + assert records["x"]["param_shapes"] == [(3,), (3,)] + + n_dim = int(compiled.n_dim) + assert n_dim == 5 + p_sl = _var_slice(compiled, p_name) + sigma_sl = _var_slice(compiled, "pop_sigma_log__") + x_sl = _var_slice(compiled, "x") + + rng = np.random.default_rng(4) + y = rng.normal(size=n_dim) + h_mu = np.array([0.6, 0.2, 0.9]) + h_sigma = np.array([0.4, 0.7, 0.1]) + flow_params = np.concatenate([h_mu, h_sigma]) + + value, ljd_c = constrain_fn(y, flow_params) + p_con = 1.0 / (1.0 + np.exp(-y[p_sl][0])) # logodds backward + sigma_con = np.exp(y[sigma_sl][0]) + expected_x = (y[x_sl] - (1 - h_mu) * p_con) * sigma_con**h_sigma + p_con + np.testing.assert_allclose(value[x_sl], expected_x, atol=1e-10) + np.testing.assert_allclose(ljd_c, (h_sigma * np.log(sigma_con)).sum(), atol=1e-10) + + y_back, ljd_u = unconstrain_fn(value, flow_params) + np.testing.assert_allclose(y_back, y, atol=1e-10) + np.testing.assert_allclose(ljd_c + ljd_u, 0.0, atol=1e-10) + + +@pytest.mark.pymc +@pytest.mark.flow +def test_build_auto_flow_roundtrip(): + import jax.numpy as jnp + + from nutpie.normalizing_flow import AutoFlow + + m, _ = _funnel_model() + compiled = nutpie.compile_pymc_model(m, backend="jax", gradient_backend="jax") + flow = build_auto_flow(m, compiled, init_params=jnp.full((5,), 0.3)) + assert isinstance(flow, AutoFlow) + assert flow.shape == (int(compiled.n_dim),) + + rng = np.random.default_rng(0) + y = jnp.asarray(rng.normal(size=flow.shape)) + x, ljd = flow.transform_and_log_det(y) + y_back, ljd_back = flow.inverse_and_log_det(x) + np.testing.assert_allclose(np.asarray(y_back), np.asarray(y), atol=1e-10) + np.testing.assert_allclose(float(ljd) + float(ljd_back), 0.0, atol=1e-10) + assert not np.isclose(float(ljd), 0.0) + + +@pytest.mark.pymc +@pytest.mark.flow +def test_auto_reparam_compile_api(capsys): + from nutpie.normalizing_flow import AutoFlow + + m, _ = _funnel_model() + compiled = nutpie.compile_pymc_model( + m, backend="jax", gradient_backend="jax", auto_reparam=True + ) + summary = capsys.readouterr().out + assert "reparametrizing 1 of 2 free variables" in summary + assert "ind_mu (AffineFlow)" in summary + auto_flow = compiled._transform_adapt_args["auto_flow"] + assert isinstance(auto_flow, AutoFlow) + + tuned = compiled.with_transform_adapt(num_layers=0) + assert tuned._transform_adapt_args["auto_flow"] is auto_flow + assert tuned._transform_adapt_args["num_layers"] == 0 + cleared = tuned.with_transform_adapt(auto_flow=None) + assert "auto_flow" not in cleared._transform_adapt_args + + with pytest.raises(ValueError, match="auto_reparam"): + nutpie.compile_pymc_model( + m, backend="jax", gradient_backend="pytensor", auto_reparam=True + ) + + +@pytest.mark.pymc +@pytest.mark.flow +def test_auto_reparam_compile_api_numba(capsys): + from nutpie.transform_adapter_numba import NumbaAutoFlow + + m, _ = _funnel_model() + compiled = nutpie.compile_pymc_model(m, backend="numba", auto_reparam=True) + summary = capsys.readouterr().out + assert "reparametrizing 1 of 2 free variables" in summary + flow = compiled._transform_adapt_args["numba_flow"] + assert isinstance(flow, NumbaAutoFlow) + assert flow.n_dim == compiled.n_dim + assert flow.n_params > 0 + + +@pytest.mark.pymc +@pytest.mark.flow +def test_numba_adapter_matches_jax_adapter(): + """The pytensor/numba hook and the jax/flowjax one compute the same + transformed logp, logdet and gradients for the same flow parameters.""" + import equinox as eqx + import jax.numpy as jnp + from flowjax import bijections + + from nutpie.transform_adapter import make_transform_adapter + from nutpie.transform_adapter_numba import make_numba_transform_adapter + + m, _ = _funnel_model() + c_jax = nutpie.compile_pymc_model( + m, backend="jax", gradient_backend="jax", auto_reparam=True + ) + c_numba = nutpie.compile_pymc_model(m, backend="numba", auto_reparam=True) + n_dim = int(c_jax.n_dim) + + rng = np.random.default_rng(0) + position = rng.normal(size=n_dim) + gradient = rng.normal(size=n_dim) + + numba_flow = c_numba._transform_adapt_args["numba_flow"] + nb = make_numba_transform_adapter(numba_flow=numba_flow)( + seed=1, position=position, gradient=gradient, chain=0 + ) + h = rng.normal(size=numba_flow.n_params) * 0.3 + loc = rng.normal(size=n_dim) * 0.1 + scale = np.exp(rng.normal(size=n_dim) * 0.2) + nb._h, nb._loc, nb._scale = h, loc, scale + + auto_flow = c_jax._transform_adapt_args["auto_flow"] + jx = make_transform_adapter(auto_flow=auto_flow)( + seed=1, + position=position, + gradient=gradient, + chain=0, + logp_fn=c_jax._raw_logp_fn, + ) + flow_jax = eqx.tree_at(lambda f: f.params, auto_flow, jnp.asarray(h)) + jx._bijection = bijections.Chain( + [bijections.Affine(jnp.asarray(loc), jnp.asarray(scale)), flow_jax] + ) + + z = rng.normal(size=n_dim) + nb_out = nb.init_from_transformed_position(z) + jx_out = jx.init_from_transformed_position(z) + for got, want in zip(nb_out, jx_out, strict=True): + np.testing.assert_allclose(got, want, rtol=1e-8, atol=1e-8) + + x = rng.normal(size=n_dim) + for got, want in zip( + nb.init_from_untransformed_position(x), + jx.init_from_untransformed_position(x), + strict=True, + ): + np.testing.assert_allclose(got, want, rtol=1e-8, atol=1e-8) + + gx = rng.normal(size=n_dim) + for got, want in zip(nb.inv_transform(x, gx), jx.inv_transform(x, gx), strict=True): + np.testing.assert_allclose(got, want, rtol=1e-8, atol=1e-8) + + +@pytest.mark.pymc +@pytest.mark.flow +def test_auto_reparam_sampling_numba(): + m, _ = _funnel_model() + compiled = nutpie.compile_pymc_model(m, backend="numba", auto_reparam=True) + trace = nutpie.sample( + compiled, chains=1, seed=1, adaptation="flow", tune=1000, draws=500 + ) + assert float(trace.sample_stats.diverging.sum()) <= 5 + np.testing.assert_allclose( + float(trace.posterior.pop_sigma_log.std()), 1.0, atol=0.3 + ) + + +@pytest.mark.pymc +@pytest.mark.flow +@pytest.mark.parametrize("n_layers", [0, 2]) +def test_auto_flow_is_outermost_bijection(n_layers): + """The VIP flow's constrain output is the model's value vector, so it + must sit at the value-space end of the chain; the diag affine and any + coupling layers operate in its base space.""" + from nutpie.normalizing_flow import AutoFlow, make_flow + + m, _ = _funnel_model() + compiled = nutpie.compile_pymc_model(m, backend="jax", gradient_backend="jax") + auto_flow = build_auto_flow(m, compiled) + + rng = np.random.default_rng(0) + n_dim = int(compiled.n_dim) + positions = rng.normal(size=(10, n_dim)) + gradients = rng.normal(size=(10, n_dim)) + chain = make_flow(1, positions, gradients, n_layers=n_layers, auto_flow=auto_flow) + assert isinstance(chain.bijections[-1], AutoFlow) + + +@pytest.mark.pymc +@pytest.mark.flow +def test_auto_reparam_nothing_found(): + with pm.Model() as m: + pm.Normal("x", 0, 1, shape=(3,)) + + with pytest.warns(UserWarning, match="did not find any variables"): + compiled = nutpie.compile_pymc_model( + m, backend="jax", gradient_backend="jax", auto_reparam=True + ) + assert "auto_flow" not in (compiled._transform_adapt_args or {}) + + +@pytest.mark.pymc +@pytest.mark.flow +def test_multiple_auto_flows_chain_to_one(): + from flowjax import bijections + + from nutpie.transform_adapter import make_transform_adapter + + m, _ = _funnel_model() + compiled = nutpie.compile_pymc_model(m, backend="jax", gradient_backend="jax") + flow = build_auto_flow(m, compiled) + adapter = make_transform_adapter(auto_flow=[flow, flow]) + chained = adapter.keywords["make_flow_fn"].keywords["auto_flow"] + assert isinstance(chained, bijections.Chain) + assert len(chained.bijections) == 2 + + +@pytest.mark.pymc +@pytest.mark.flow +def test_auto_reparam_sampling(): + m, _ = _funnel_model() + compiled = nutpie.compile_pymc_model( + m, backend="jax", gradient_backend="jax", auto_reparam=True + ) + trace = nutpie.sample( + compiled, chains=1, seed=1, adaptation="flow", tune=1000, draws=500 + ) + assert float(trace.sample_stats.diverging.sum()) <= 5 + np.testing.assert_allclose( + float(trace.posterior.pop_sigma_log.std()), 1.0, atol=0.3 + )