Reparameterisation

Site-level reparameterisation strategies.

reparam

Reparameterisation strategies.

Reparam is the abstract strategy interface; ReparamOrchestrator (returned by the reparam(strategies) factory) is the effect handler that dispatches each covered site to the matching strategy. The concrete strategies are LocScaleReparam (non-centred Normal), TransformReparam (fixed bijector), NeuTraReparam (warp geometry through a trained autoguide), and ConjugateReparam (analytic marginalisation for conjugate pairs).

See the individual module docstrings for citations to the original literature: Betancourt and Girolami 2015 for the non-centred rewrite, Hoffman et al. 2019 for NeuTra, and Bernardo and Smith 1994 for the conjugate pairs.

Reparam

Bases: ABC

Base class for site-level reparameterisation strategies.

A strategy is a policy, not a handler. The ReparamOrchestrator handler dispatches each covered site to the matching strategy and installs the strategy's value and density on the site. Subclasses override :meth:apply.

apply abstractmethod

apply(site: SiteRequest) -> tuple[Tensor, Tensor]

Rewrite one site.

PARAMETER	DESCRIPTION
site	The site, with a value fixed when one is given. TYPE: SiteRequest

RETURNS	DESCRIPTION
tuple[Tensor, Tensor]	The value the site takes, in the site's original space so downstream sites see the intended value, and the density a downstream sampler should score the site with.

Source code in src/quivers/effects/reparam/base.py

@abstractmethod
def apply(self, site: SiteRequest) -> tuple[torch.Tensor, torch.Tensor]:
    """Rewrite one site.

    Parameters
    ----------
    site : SiteRequest
        The site, with a value fixed when one is given.

    Returns
    -------
    tuple[torch.Tensor, torch.Tensor]
        The value the site takes, in the site's original space so
        downstream sites see the intended value, and the density a
        downstream sampler should score the site with.
    """

ReparamOrchestrator

ReparamOrchestrator(strategies: dict[str, Reparam])

Bases: EffectHandler

Dispatch sample sites to per-site reparam strategies.

PARAMETER	DESCRIPTION
strategies	Site name to strategy. TYPE: dict[str, Reparam]

Source code in src/quivers/effects/reparam/base.py

def __init__(self, strategies: dict[str, Reparam]) -> None:
    self.strategies = dict(strategies)

install

install(run: RunContext) -> tuple[Installation, ...]

Install the dispatching handler of the random instance.

PARAMETER	DESCRIPTION
run	The run being prepared. TYPE: RunContext

RETURNS	DESCRIPTION
tuple[Installation, ...]	The handler, which answers covered sites and forwards the rest.

Source code in src/quivers/effects/reparam/base.py

def install(self, run: RunContext) -> tuple[Installation, ...]:
    """Install the dispatching handler of the ``random`` instance.

    Parameters
    ----------
    run : RunContext
        The run being prepared.

    Returns
    -------
    tuple[Installation, ...]
        The handler, which answers covered sites and forwards the
        rest.
    """
    kernel = run.kernel
    score_instance = kernel.score.entry.instance
    seen: dict[str, TorchSampleable] = {}
    definition = _handler_def(
        "Random.reparam",
        RANDOM,
        ((RANDOM_SAMPLE, ResumptionGrade.LINEAR),),
        key=f"reparam-{id(self):x}",
        input_type=run.result_type,
        output_type=run.result_type,
        introduced=EffectRow((RowEntry(score_instance, SCORE),)),
        total=False,
    )

    def sample(
        request: RuntimeRequest, resume: Resumption, context: ClauseContext
    ) -> object:
        """Answer a covered site with its strategy's value and density.

        Parameters
        ----------
        request : RuntimeRequest
            The request being answered and its arguments.
        resume : Resumption
            The continuation.
        context : ClauseContext
            Runtime services available to the clause.

        Returns
        -------
        object
            What the resumed computation produced, or a `Forward`
            for a site no strategy covers.
        """
        site, sampleable = _site_and_sampleable(request, definition.name)
        assert isinstance(site, str)
        assert isinstance(sampleable, TorchSampleable)
        strategy = self.strategies.get(site)
        if strategy is None:
            seen[site] = sampleable
            return Forward()
        value, log_prob = strategy.apply(
            SiteRequest(
                site,
                sampleable,
                seen=dict(seen),
                observation=run.observations.get(site),
            )
        )
        seen[site] = sampleable
        _emit_score(
            context,
            request,
            score_instance,
            log_prob,
            run.validator,
            role="draw-score",
        )
        return resume(value)

    runtime = RuntimeHandler(
        definition, {RANDOM_SAMPLE: RuntimeClause(sample, run.validator)}
    )
    return (Installation(kernel.random, runtime),)

SiteRequest dataclass

SiteRequest(name: str, sampleable: TorchSampleable, given: Tensor | None = None, seen: Mapping[str, TorchSampleable] = None, observation: Tensor | None = None)

What a strategy is given for one sample site.

PARAMETER	DESCRIPTION
name	The site's name. TYPE: str
sampleable	The site's distribution at its input. TYPE: TorchSampleable
given	A value the site is to take, when an enclosing strategy or a caller fixed one; None when the strategy draws. TYPE: Tensor | None DEFAULT: None
seen	The sampleables of the sites the orchestrator saw before this one, by name, so a strategy may read a parent's parameters. TYPE: Mapping[str, TorchSampleable] DEFAULT: None
observation	The run's observation for the site, if any. TYPE: Tensor | None DEFAULT: None

with_value

with_value(value: Tensor) -> SiteRequest

The same request with a value fixed.

PARAMETER	DESCRIPTION
value	The value the site is to take. TYPE: Tensor

RETURNS	DESCRIPTION
SiteRequest	The request with given set.

Source code in src/quivers/effects/reparam/base.py

def with_value(self, value: torch.Tensor) -> SiteRequest:
    """The same request with a value fixed.

    Parameters
    ----------
    value : torch.Tensor
        The value the site is to take.

    Returns
    -------
    SiteRequest
        The request with ``given`` set.
    """
    return replace(self, given=value)

ConjugateReparam

ConjugateReparam(parent: str, parent_family: str, child_family: str)

Bases: Reparam

Score a child site under the conjugate marginal of its parent.

Applied to the child site of a conjugate pair, the strategy answers the site with the run's observation for it, or a draw from the child's own sampleable when there is none, and scores it under the closed-form marginal that integrates the parent out, read from the parent's parameters as the orchestrator saw them. The parent's own density is left in the joint; collapse the parent with quivers.effects.collapse when the marginal joint is wanted.

PARAMETER	DESCRIPTION
parent	The parent site's name. TYPE: str
parent_family	The parent's family. TYPE: str
child_family	The child's family. TYPE: str

RAISES	DESCRIPTION
KeyError	If the pair has no registered marginal.

Source code in src/quivers/effects/reparam/conjugate.py

def __init__(self, parent: str, parent_family: str, child_family: str) -> None:
    self.parent = parent
    self.parent_family = parent_family
    self.child_family = child_family
    self.solver = conjugate_solver(parent_family, child_family)

apply

apply(site: SiteRequest) -> tuple[Tensor, Tensor]

Rewrite the child site.

PARAMETER	DESCRIPTION
site	The child site. TYPE: SiteRequest

RETURNS	DESCRIPTION
tuple[Tensor, Tensor]	The value and its marginal log density.

RAISES	DESCRIPTION
KeyError	If the orchestrator has not seen the parent site.
ValueError	If the parent or the child is not of the declared family.

Source code in src/quivers/effects/reparam/conjugate.py

def apply(self, site: SiteRequest) -> tuple[torch.Tensor, torch.Tensor]:
    """Rewrite the child site.

    Parameters
    ----------
    site : SiteRequest
        The child site.

    Returns
    -------
    tuple[torch.Tensor, torch.Tensor]
        The value and its marginal log density.

    Raises
    ------
    KeyError
        If the orchestrator has not seen the parent site.
    ValueError
        If the parent or the child is not of the declared family.
    """
    try:
        parent = site.seen[self.parent]
    except KeyError as error:
        raise KeyError(
            f"ConjugateReparam: parent site {self.parent!r} was not seen "
            f"before child {site.name!r}"
        ) from error
    if family_of(parent) != self.parent_family:
        raise ValueError(
            f"ConjugateReparam: parent {self.parent!r} draws from "
            f"{family_of(parent)!r}, not {self.parent_family!r}"
        )
    if family_of(site.sampleable) != self.child_family:
        raise ValueError(
            f"ConjugateReparam: child {site.name!r} draws from "
            f"{family_of(site.sampleable)!r}, not {self.child_family!r}"
        )
    if site.given is not None:
        value = site.given
    elif site.observation is not None:
        value = site.observation
    else:
        value = site.sampleable.rsample()
    parent_value = site.sampleable.input
    return value, self.solver(
        parent.parameters(), parent_value, site.sampleable, value
    )

LocScaleReparam

LocScaleReparam(centered: float = 0.0)

Bases: Reparam

Non-centred rewrite for location-scale sample sites.

The site's morphism must expose a _get_params(x) method returning (loc, scale) tensors of the same shape, the convention every ConditionalNormal-shaped family in quivers.continuous.families follows.

PARAMETER	DESCRIPTION
centered	Interpolation between fully centred (1.0) and fully non-centred (0.0) parameterisations. Values in between produce a partial reparam, matching the centered parameter of Pyro's LocScaleReparam. TYPE: float DEFAULT: 0.0

RAISES	DESCRIPTION
ValueError	If centered is outside [0, 1].

Source code in src/quivers/effects/reparam/loc_scale.py

def __init__(self, centered: float = 0.0) -> None:
    if not 0.0 <= centered <= 1.0:
        raise ValueError(
            f"LocScaleReparam: centered must be in [0, 1], got {centered}."
        )
    self.centered = float(centered)

apply

apply(site: SiteRequest) -> tuple[Tensor, Tensor]

Rewrite one site.

PARAMETER	DESCRIPTION
site	The site. TYPE: SiteRequest

RETURNS	DESCRIPTION
tuple[Tensor, Tensor]	The value and its log density under Normal(loc, scale), summed over the trailing feature axis.

RAISES	DESCRIPTION
TypeError	If the site's morphism exposes no _get_params.

Source code in src/quivers/effects/reparam/loc_scale.py

def apply(self, site: SiteRequest) -> tuple[torch.Tensor, torch.Tensor]:
    """Rewrite one site.

    Parameters
    ----------
    site : SiteRequest
        The site.

    Returns
    -------
    tuple[torch.Tensor, torch.Tensor]
        The value and its log density under ``Normal(loc, scale)``,
        summed over the trailing feature axis.

    Raises
    ------
    TypeError
        If the site's morphism exposes no ``_get_params``.
    """
    morphism = site.sampleable.morphism
    get_params = getattr(morphism, "_get_params", None)
    if get_params is None:
        raise TypeError(
            f"LocScaleReparam: site '{site.name}' morphism "
            f"{type(morphism).__name__} does not expose `_get_params(x)`; "
            f"LocScaleReparam requires a Normal-family morphism."
        )
    loc, scale = get_params(site.sampleable.input)
    # Partial centring interpolates the effective scale used to
    # push the base sample forward: ``centered=1`` recovers the
    # original sample-then-score path, ``centered=0`` the fully
    # non-centred rewrite.
    effective = scale.pow(self.centered)
    base_scale = scale / effective
    if site.given is None:
        eps = torch.randn_like(loc)
        y_raw = effective * eps
        y = loc + base_scale * y_raw
    else:
        y = site.given
    residual = (y - loc) / scale
    log_p = -0.5 * residual.pow(2) - scale.log() - 0.5 * math.log(2 * math.pi)
    return y, log_p.sum(dim=-1)

NeuTraReparam

NeuTraReparam(autoguide: Guide)

Bases: Reparam

Reparameterise a site through a trained normalizing-flow guide.

PARAMETER	DESCRIPTION
autoguide	A trained AutoIAFGuide, or any Guide whose sample returns per-site values, whose draw at apply time has the site's support. TYPE: Guide

Source code in src/quivers/effects/reparam/neutra.py

def __init__(self, autoguide: Guide) -> None:
    self.autoguide = autoguide

apply

apply(site: SiteRequest) -> tuple[Tensor, Tensor]

Rewrite one site.

PARAMETER	DESCRIPTION
site	The site. TYPE: SiteRequest

RETURNS	DESCRIPTION
tuple[Tensor, Tensor]	The guide's draw for the site, or the site's own draw when the guide does not cover it, and its log density under the site's sampleable; scoring uses the original sampleable throughout, so a partially covered NeuTra still yields a well-defined joint.

RAISES	DESCRIPTION
KeyError	If the guide covers the site but its draw does not carry it.

Source code in src/quivers/effects/reparam/neutra.py

def apply(self, site: SiteRequest) -> tuple[torch.Tensor, torch.Tensor]:
    """Rewrite one site.

    Parameters
    ----------
    site : SiteRequest
        The site.

    Returns
    -------
    tuple[torch.Tensor, torch.Tensor]
        The guide's draw for the site, or the site's own draw when
        the guide does not cover it, and its log density under the
        site's sampleable; scoring uses the original sampleable
        throughout, so a partially covered NeuTra still yields a
        well-defined joint.

    Raises
    ------
    KeyError
        If the guide covers the site but its draw does not carry it.
    """
    registry = getattr(self.autoguide, "registry", None)
    names = set(registry.names()) if registry is not None else set()
    if site.given is not None:
        y = site.given
    elif site.name not in names:
        y = site.sampleable.rsample()
    else:
        samples = self.autoguide.sample(site.sampleable.input)
        if site.name not in samples:
            raise KeyError(
                f"NeuTraReparam: autoguide sample did not contain site "
                f"'{site.name}'."
            )
        y = samples[site.name]
    return y, site.sampleable.log_prob(y)

TransformReparam

TransformReparam(bijector: Bijector)

Bases: Reparam

Reparameterise a sample site through a fixed bijector.

The reparameterised base sample lives in the bijector's domain and the site's value in its codomain. The joint density remains that of the original distribution: the strategy scores the value under the site's own sampleable, and the bijector reshapes the sampling geometry a downstream sampler works in.

PARAMETER	DESCRIPTION
bijector	The measurable bijection b the site's geometry is reshaped through. TYPE: Bijector

Source code in src/quivers/effects/reparam/transform.py

def __init__(self, bijector: Bijector) -> None:
    self.bijector = bijector

apply

apply(site: SiteRequest) -> tuple[Tensor, Tensor]

Rewrite one site.

PARAMETER	DESCRIPTION
site	The site. TYPE: SiteRequest

RETURNS	DESCRIPTION
tuple[Tensor, Tensor]	The value, drawn from the site's sampleable unless given, and its log density under the sampleable.

Source code in src/quivers/effects/reparam/transform.py

def apply(self, site: SiteRequest) -> tuple[torch.Tensor, torch.Tensor]:
    """Rewrite one site.

    Parameters
    ----------
    site : SiteRequest
        The site.

    Returns
    -------
    tuple[torch.Tensor, torch.Tensor]
        The value, drawn from the site's sampleable unless given, and
        its log density under the sampleable.
    """
    y = site.given if site.given is not None else site.sampleable.rsample()
    return y, site.sampleable.log_prob(y)

reparam

reparam(strategies: dict[str, Reparam]) -> ReparamOrchestrator

Return a ReparamOrchestrator for the given per-site strategies.

PARAMETER	DESCRIPTION
strategies	Site name to strategy. TYPE: dict[str, Reparam]

RETURNS	DESCRIPTION
ReparamOrchestrator	The handler.

Source code in src/quivers/effects/reparam/base.py

def reparam(strategies: dict[str, Reparam]) -> ReparamOrchestrator:
    """Return a `ReparamOrchestrator` for the given per-site strategies.

    Parameters
    ----------
    strategies : dict[str, Reparam]
        Site name to strategy.

    Returns
    -------
    ReparamOrchestrator
        The handler.
    """
    return ReparamOrchestrator(strategies)