Probabilistic

class pmrf.models.composite.wrapped.Probabilistic(model: ~pmrf.models.base.Model, distribution: ~distreqx.distributions._distribution.AbstractDistribution, target: ~typing.Callable[[~typing.Any], ~typing.Any] = <function Probabilistic.<lambda>>, constraint: ~parax.constraints.AbstractConstraint | None = None, static: ~typing.Any = None)

Bases: Model, AbstractUnwrappable

(experimental) A wrapper to make an existing model probabilistic.

This provides the ability to associate a probability distribution with a model or one of its sub-models/parameters after the model was created.

This is a useful for advanced use-cases where you want to attach a distribution to an entire model (perhaps overriding previous distributions on lower levels), as opposed to more standard cases where you want to model the distributions of individual variables (in which case you should likely use pmrf.Random instead).

Variables:

probabilistic (Model | parax.Probabilize | parax.Combine) – The updated structure containing the parax.Probabilize node.

Examples

>>> from pmrf.models import Probabilistic, Resistor
>>> from pmrf.distributions import Normal, Joint
>>>
>>> res = Resistor(R=50.0)
>>>
>>> # Use Case 1: Target a specific parameter (leaf)
>>> prob_res_leaf = Probabilistic(
...     model=res,
...     distribution=Normal(loc=50.0, scale=1.0),
...     target=lambda m: m.R
... )
>>>
>>> # Use Case 2: Wrap the entire model (requires matching distribution tree)
>>> import equinox as eqx
>>> dist_tree = Joint(eqx.tree_at(lambda m: m.R, res, Normal(loc=50.0, scale=1.0)))
>>> prob_res_tree = Probabilistic(
...     model=res,
...     distribution=dist_tree,
... )

Initialize the Probabilistic model.

Parameters:

  • model (Model) – The base model to wrap.

  • distribution (AbstractDistribution or tuple of AbstractDistribution) – The probability distribution(s) to associate with the target(s).

  • target (callable, optional) – A callable (lens) extracting the parameter(s) to make probabilistic (e.g., lambda m: m.R or lambda m: (m.R, m.C)). Defaults to the identity function.

  • constraint (AbstractConstraint or tuple of AbstractConstraint, optional) – Optional constraint(s) for the distribution(s).

  • static (Any or tuple of Any, optional) – A pytree structurally complementary to target’s value (None at every leaf the target owns), left untouched and recombined via parax.Combine. Use this when target only covers part of a model (e.g. the leaves a trained flow was fit on) and the rest should keep its own existing parameters/priors as-is, rather than needing the distribution to cover the whole model.

unwrap() → Model

Implements parax.AbstractUnwrappable so nested Probabilistic nodes fully collapse too.

property model: Model

Returns the underlying model (or the probabilistic wrapper if applied to the root).

probabilistic: Probabilize | Combine | Model