Repository: liesel-devs/liesel
Pull request #253: Future-proof MCMC convenience
Author: login: jobrachem
Labels: name: comp:goose
description: This issue is related to the goose module
name: comp:model
description: This issue is related to the model module
Created: 2025-04-04T20:33:59Z | Merged: 2025-04-28T16:08:12Z

Turns out, I also really like this way of doing it. Things now revolve around the following arguments and objects:

- `gs.MCMCSpec(kernel, kernel_kwargs, kernel_group, jitter_dist)` holds information about MCMC inference for a single variable, but allows for passing a `kernel_group` argument that is used to identify kernels that cover multiple variables.
- `gs.LieselMCMC(model: lsl.Model)` collects MCMC inference information from variables in the model
- `lsl.Var(inference)` new argument that accepts instances or dicts of `MCMCSpec`.

Here's a code example that showcases lots of functionality already:

```python
mu = lsl.Var.new_param(
    0.0,
    lsl.Dist(tfd.Normal, loc=0.0, scale=1.0),
    inference=gs.MCMCSpec(
        gs.NUTSKernel,
        kernel_group="a",
        jitter_dist=tfd.Uniform(low=-1.0, high=1.0),
    ),
    name="mu",
)

sigma = lsl.Var.new_param(
    1.0,
    lsl.Dist(tfd.InverseGamma, concentration=1.0, scale=0.5),
    inference=gs.MCMCSpec(
        gs.NUTSKernel,
        kernel_group="a",
        jitter_dist=tfd.Uniform(low=0.0, high=1.0),
    ),
    name="sigma",
)

model = lsl.Model([mu, sigma])

mcmc = gs.LieselMCMC(model)

mcmc.kernel_list() # list of mcmc kernels
mcmc.jitter_functions() # dictionary of jitter functions
mcmc.engine_builder( # shortcut for an engine builder
    seed=1,
    num_chains=4,
    warmup_duration=200,
    posterior_duration=100,
    term_duration=50, # from here on, the arguments have defaults
    thinning_warmup=1,
    thinning_posterior=1,
    apply_jitter=True
) 
mcmc.engine() # like .engine_builder(), but directly builds the engine
```

From here, it was quite easy to allow multiple specs:

```python
mu = lsl.Var.new_param(
    0.0,
    lsl.Dist(tfd.Normal, loc=0.0, scale=1.0),
    inference={
        "a": gs.MCMCSpec(gs.NUTSKernel),
        "b": gs.MCMCSpec(gs.IWLSKernel),
    },
    name="mu",
)

sigma = lsl.Var.new_param(
    1.0,
    lsl.Dist(tfd.InverseGamma, concentration=2.0, scale=1.0),
    inference=gs.MCMCSpec(gs.IWLSKernel),
    name="sigma",
)

# will take spec 'a' from mu
# and the single spec from sigma
engine = gs.LieselMCMC(model, which="a").engine()
```



Diff:
diff --git a/liesel/goose/__init__.py b/liesel/goose/__init__.py
index b44bc1d5..0d32ccca 100644
--- a/liesel/goose/__init__.py
+++ b/liesel/goose/__init__.py
@@ -15,6 +15,8 @@
     NamedTupleInterface,
 )
 from .iwls import IWLSKernel
+from .mcmc_spec import LieselMCMC as LieselMCMC
+from .mcmc_spec import MCMCSpec as MCMCSpec
 from .mh_kernel import MHKernel, MHProposal
 from .models import DataClassModel, DictModel
 from .nuts import NUTSKernel
@@ -40,6 +42,8 @@
     "LieselInterface",
     "Engine",
     "EngineBuilder",
+    "MCMCSpec",
+    "LieselMCMC",
     "EpochConfig",
     "EpochType",
     "GibbsKernel",
diff --git a/liesel/goose/builder.py b/liesel/goose/builder.py
index d71f4620..cb7e9ebd 100644
--- a/liesel/goose/builder.py
+++ b/liesel/goose/builder.py
@@ -7,6 +7,8 @@
 implementations.
 """
 
+from __future__ import annotations
+
 import logging
 import math
 from collections.abc import Iterable
diff --git a/liesel/goose/mcmc_spec.py b/liesel/goose/mcmc_spec.py
new file mode 100644
index 00000000..736168fa
--- /dev/null
+++ b/liesel/goose/mcmc_spec.py
@@ -0,0 +1,254 @@
+from __future__ import annotations
+
+import logging
+from collections.abc import Callable
+from dataclasses import dataclass, field
+from typing import TYPE_CHECKING, Any
+
+import jax.numpy as jnp
+import tensorflow_probability.substrates.jax.distributions as tfd
+
+from .builder import EngineBuilder
+from .interface import LieselInterface
+from .types import Array, JitterFunctions, Kernel, KeyArray
+
+if TYPE_CHECKING:
+    from liesel.model import Model, Var
+
+
+logger = logging.getLogger(__name__)
+
+
+@dataclass
+class LieselMCMC:
+    """
+    Manages the setup of MCMC specifications for a Liesel model.
+
+    Parameters
+    ----------
+    model
+        The Liesel model object containing the variables and their inference \
+        specifications.
+    which
+        A named inference configuration to use. If None, the default inference \
+        attached to each variable is used.
+    """
+
+    model: Model
+    which: str | None = None
+
+    def get_spec(self, var: Var) -> MCMCSpec | None:
+        """
+        Retrieve the MCMC specification for a given variable.
+
+        Parameters
+        ----------
+        var
+            The model variable for which to get the MCMC specification.
+
+        Returns
+        -------
+        The MCMC specification if available, otherwise None.
+
+        Raises
+        ------
+        ValueError
+            If the inference attached to the variable is not of type ``MCMCSpec``.
+        """
+        inference = var.get_inference(self.which)
+        if inference is None:
+            return inference
+
+        if not isinstance(inference, MCMCSpec):
+            raise ValueError(
+                f"Attribute 'inference' of variable {var} is of type"
+                f" {type(inference)}, but expected type '{MCMCSpec}'."
+            )
+        return inference
+
+    def get_kernel_groups(self) -> dict[str, _KernelGroup]:
+        """
+        Collect and organize model variables into kernel groups.
+
+        Returns
+        -------
+        A dictionary mapping group names or variable names to their corresponding \
+        kernel group specifications.
+
+        Raises
+        ------
+        ValueError
+            If variables in the same kernel group have inconsistent kernels or kernel \
+            arguments.
+        """
+        vars_ = self.model.vars
+        kernel_groups: dict[str, _KernelGroup] = {}
+
+        for name, var in vars_.items():
+            inference = self.get_spec(var)
+
+            if not inference:
+                continue
+
+            group_name = inference.kernel_group
+
+            if group_name is None:
+                kernel_groups[name] = _KernelGroup(
+                    kernel=inference.kernel,
+                    kwargs=inference.kernel_kwargs,
+                    position_keys=[name],
+                )
+
+            elif group_name in kernel_groups:
+                group = kernel_groups[group_name]
+                same_kernel = group.kernel is inference.kernel
+                if not same_kernel:
+                    raise ValueError(
+                        "Found incoherent kernel classes for kernel group"
+                        f" {group_name}."
+                    )
+
+                same_kwargs = group.kwargs == inference.kernel_kwargs
+                if not same_kwargs:
+                    raise ValueError(
+                        f"Found incoherent kwargs for kernel group {group_name}."
+                    )
+                group.position_keys.append(name)
+
+            else:
+                kernel_groups[group_name] = _KernelGroup(
+                    kernel=inference.kernel,
+                    kwargs=inference.kernel_kwargs,
+                    position_keys=[name],
+                )
+
+        return kernel_groups
+
+    def get_kernel_list(self) -> list[Kernel]:
+        """
+        Construct the list of MCMC kernels from kernel groups.
+
+        Returns
+        -------
+        A list of initialized kernel instances ready to be added to the MCMC engine.
+        """
+        kernel_groups = self.get_kernel_groups()
+        kernel_list = [
+            g.kernel(g.position_keys, **g.kwargs)  # type: ignore
+            for g in kernel_groups.values()
+        ]
+        return kernel_list
+
+    def get_jitter_functions(self) -> JitterFunctions:
+        """
+        Collect jitter functions for model variables that define a jitter distribution.
+
+        Returns
+        -------
+        A dictionary mapping variable names to their jitter application functions.
+        """
+        jitter_functions: JitterFunctions = {}
+        for name, var in self.model.vars.items():
+            inference = self.get_spec(var)
+            if inference is not None and inference.jitter_dist is not None:
+                jitter_functions[name] = inference.apply_jitter
+
+        return jitter_functions
+
+    def get_engine_builder(
+        self,
+        seed: int,
+        num_chains: int,
+        apply_jitter: bool = True,
+    ) -> EngineBuilder:
+        """
+        Create and configure an `EngineBuilder` for MCMC sampling.
+
+        Parameters
+        ----------
+        seed
+            Random seed for reproducibility.
+        num_chains
+            Number of MCMC chains.
+        apply_jitter
+            Whether to apply jitter to the initial states, by default True.
+
+        Returns
+        -------
+        EngineBuilder
+            A configured ``EngineBuilder`` instance.
+        """
+        eb = EngineBuilder(seed=seed, num_chains=num_chains)
+        eb.set_model(LieselInterface(self.model))
+        eb.set_initial_values(self.model.state)
+
+        for kernel in self.get_kernel_list():
+            eb.add_kernel(kernel)
+
+        if apply_jitter:
+            eb.set_jitter_fns(self.get_jitter_functions())
+
+        return eb
+
+
+@dataclass
+class _KernelGroup:
+    kernel: Callable[..., Kernel]
+    kwargs: dict[str, Any] = field(default_factory=dict)
+    position_keys: list[str] = field(default_factory=list)
+
+
+@dataclass
+class MCMCSpec:
+    """
+    Specification for the MCMC kernel and optional jitter distribution associated with a
+    model variable.
+
+    Parameters
+    ----------
+    kernel
+        A callable that returns a ``Kernel`` instance when provided with position keys \
+        and keyword arguments.
+    kernel_kwargs
+        Additional keyword arguments to be passed to the kernel callable.
+    kernel_group
+        Name of the kernel group this variable belongs to. Variables in the same group \
+        must share the same kernel type and arguments.
+    jitter_dist
+        A TensorFlow Probability distribution used to apply random jitter to the \
+        initial value of the variable.
+    """
+
+    kernel: Callable[..., Kernel]
+    kernel_kwargs: dict[str, Any] = field(default_factory=dict)
+    kernel_group: str | None = None
+    jitter_dist: tfd.Distribution | None = None
+
+    def apply_jitter(self, seed: KeyArray, value: Array) -> Array:
+        """
+        Apply random jitter to a given value using the specified jitter distribution.
+
+        If a jitter distribution is set, a random sample from the distribution is added
+        to the original value. If no jitter distribution is set, the original value is
+        returned unchanged.
+
+        Parameters
+        ----------
+        seed
+            A PRNG key used for random sampling.
+        value
+            The value to which jitter should be applied.
+
+        Returns
+        -------
+        The jittered value with the same shape as the input.
+        """
+        if self.jitter_dist is None:
+            return value
+
+        original_shape = jnp.asarray(value).shape
+        value = jnp.atleast_1d(value)
+
+        jitter = self.jitter_dist.sample(sample_shape=jnp.shape(value), seed=seed)
+
+        return jnp.reshape(value + jitter, original_shape)
diff --git a/liesel/model/nodes.py b/liesel/model/nodes.py
index e18ab0c2..754a824f 100644
--- a/liesel/model/nodes.py
+++ b/liesel/model/nodes.py
@@ -32,8 +32,11 @@
 from .viz import plot_nodes, plot_vars
 
 if TYPE_CHECKING:
+    from ..goose import MCMCSpec
     from .model import Model
 
+    InferenceTypes = Union[MCMCSpec, None, dict[str, MCMCSpec], Any]
+
 __all__ = [
     "Array",
     "Bijector",
@@ -1160,6 +1163,8 @@ class Var:
     name
         The name of the variable. If you do not specify a name, a unique name will be \
         automatically generated upon initialization of a :class:`.Model`.
+    inference
+        Additional information that can be used to set up inference algorithms.
 
     See Also
     --------
@@ -1189,6 +1194,7 @@ class Var:
 
     __slots__ = (
         "info",
+        "inference",
         "_auto_transform",
         "_dist_node",
         "_groups",
@@ -1205,6 +1211,7 @@ def __init__(
         value: Any,
         distribution: Dist | None = None,
         name: str = "",
+        inference: InferenceTypes = None,
     ):
         self._name = name
         self._value_node: Node = Value(None)
@@ -1229,9 +1236,15 @@ def __init__(
         self.info: dict[str, Any] = {}
         """Additional meta-information about the variable as a dict."""
 
+        self.inference = inference
+
     @classmethod
     def new_param(
-        cls, value: Any, distribution: Dist | None = None, name: str = ""
+        cls,
+        value: Any,
+        distribution: Dist | None = None,
+        name: str = "",
+        inference: InferenceTypes = None,
     ) -> Var:
         """
         Initializes a strong variable that acts as a model parameter.
@@ -1249,6 +1262,8 @@ def new_param(
         name
             The name of the variable. If you do not specify a name, a unique name will \
             be automatically generated upon initialization of a :class:`.Model`.
+        inference
+            Additional information that can be used to set up inference algorithms.
 
         See Also
         --------
@@ -1274,7 +1289,7 @@ def new_param(
         Var(name="")
 
         """
-        var = cls(value, distribution, name)
+        var = cls(value, distribution, name, inference=inference)
         var.value_node.monitor = True
         var.parameter = True
         return var
@@ -1433,7 +1448,9 @@ def new_calc(
         return var
 
     @classmethod
-    def new_value(cls, value: Any, name: str = "") -> Var:
+    def new_value(
+        cls, value: Any, name: str = "", inference: InferenceTypes = None
+    ) -> Var:
         """
         Initializes a strong variable without a distribution.
 
@@ -1446,6 +1463,8 @@ def new_value(cls, value: Any, name: str = "") -> Var:
         name
             The name of the variable. If you do not specify a name, a unique name will \
             be automatically generated upon initialization of a :class:`.Model`.
+        inference
+            Additional information that can be used to set up inference algorithms.
 
         See Also
         --------
@@ -1464,9 +1483,18 @@ def new_value(cls, value: Any, name: str = "") -> Var:
         Var(name="")
 
         """
-        var = cls(value, name=name)
+        var = cls(value, name=name, inference=inference)
         return var
 
+    def get_inference(self, key: str | None) -> InferenceTypes:
+        if isinstance(self.inference, dict):
+            if key is None:
+                raise ValueError(
+                    f"{key=} is invalid. Possible keys: {list(self.inference)}."
+                )
+            return self.inference[key]
+        return self.inference
+
     def all_input_nodes(self) -> tuple[Node, ...]:
         """Returns all input *nodes* as a unique tuple."""
         inputs1 = self.value_node.all_input_nodes()
@@ -1501,6 +1529,7 @@ def transform(
         self,
         bijector: type[jb.Bijector] | jb.Bijector | None = None,
         *bijector_args,
+        inference: InferenceTypes | Literal["drop"] = None,
         name: str | None = None,
         **bijector_kwargs,
     ) -> Var:
@@ -1525,6 +1554,14 @@ def transform(
             directly.
         bijector_args
             The arguments passed on to the init function of the bijector.
+        inference
+            Additional information that can be used to set up inference algorithms for \
+            the new, transformed variable. If ``"drop"``, the inference \
+            information will be dropped from the original variable. \
+            The new variable will have no inference information. \
+            If ``None`` (default), the new variable will likewise have no inference \
+            information, but an error will be raised if there is inference information \
+            on the original variable.
         name
             Name for the new, transformed variable. If ``None`` (default), the new \
             name will be ``<old_name>_transformed``, where ``<old_name>`` is \
@@ -1542,7 +1579,10 @@ def transform(
             If the variable is weak or if the variable has no distribution.
         ValueError
             If the argument ``bijector`` is ``None``, but the distribution does
-            not have a default event space bijector.
+            not have a default event space bijector. Also, if in the arguments to
+            :meth:`.transform` is ``inference=None`` but the variable
+            attribute :attr:`.inference` is not ``None``.
+
 
         Notes
         -----
@@ -1615,6 +1655,12 @@ def transform(original_var: lsl.Var, bijector: tfb.Bijector):
         >>> scale.update().log_prob
         0.0
         """
+        if inference is None and self.inference:
+            raise ValueError(
+                f"{self} has inference information in the .inference attribute. "
+                "To proceed with transformation, the .inference information needs to"
+                "be explicitly removed. You can transform with ``inference='drop'``."
+            )
         # if self.weak:
         #     raise RuntimeError(f"{repr(self)} is weak")
 
@@ -1679,6 +1725,11 @@ def transform(original_var: lsl.Var, bijector: tfb.Bijector):
         if name is not None:
             tvar.name = name
 
+        if inference == "drop":
+            self.inference = None
+        else:
+            self.inference = None
+            tvar.inference = inference
         return tvar
 
     @in_model_method
diff --git a/tests/goose/test_builder.py b/tests/goose/test_builder.py
index 85b09353..0d07fa42 100644
--- a/tests/goose/test_builder.py
+++ b/tests/goose/test_builder.py
@@ -4,9 +4,12 @@
 
 import jax
 import jax.numpy as jnp
+import pytest
+import tensorflow_probability.substrates.jax.bijectors as tfb
 import tensorflow_probability.substrates.jax.distributions as tfd
 
 import liesel.goose as gs
+import liesel.model as lsl
 from liesel.goose.builder import EngineBuilder
 from liesel.goose.interface import DictInterface
 
@@ -51,3 +54,239 @@ def test_jitter_fns():
     assert not jnp.allclose(ms["y"], engine._model_states["y"][1])
 
     assert not jnp.allclose(engine._model_states["x"][0], engine._model_states["x"][1])
+
+
+class TestLieselMCMC:
+    def test_engine(self):
+        mu = lsl.Var.new_param(
+            0.0,
+            lsl.Dist(tfd.Normal, loc=0.0, scale=1.0),
+            inference=gs.MCMCSpec(gs.NUTSKernel),
+            name="mu",
+        )
+
+        model = lsl.Model([mu])
+
+        mcmc = gs.LieselMCMC(model)
+        eb = mcmc.get_engine_builder(seed=1, num_chains=4)
+        eb.set_duration(warmup_duration=200, posterior_duration=100)
+        engine = eb.build()
+
+        engine.sample_all_epochs()
+        results = engine.get_results()
+        samples = results.get_posterior_samples()
+
+        assert "mu" in samples
+
+    def test_multiple_specs(self):
+        mu = lsl.Var.new_param(
+            0.0,
+            lsl.Dist(tfd.Normal, loc=0.0, scale=1.0),
+            inference={
+                "a": gs.MCMCSpec(gs.NUTSKernel),
+                "b": gs.MCMCSpec(gs.IWLSKernel),
+            },
+            name="mu",
+        )
+
+        model = lsl.Model([mu])
+
+        with pytest.raises(ValueError):
+            gs.LieselMCMC(model).get_kernel_list()
+
+        kernels = gs.LieselMCMC(model, which="a").get_kernel_list()
+        assert isinstance(kernels[0], gs.NUTSKernel)
+
+        kernels = gs.LieselMCMC(model, which="b").get_kernel_list()
+        assert isinstance(kernels[0], gs.IWLSKernel)
+
+    def test_multiple_and_single_specs(self):
+        mu = lsl.Var.new_param(
+            0.0,
+            lsl.Dist(tfd.Normal, loc=0.0, scale=1.0),
+            inference={
+                "a": gs.MCMCSpec(gs.NUTSKernel),
+                "b": gs.MCMCSpec(gs.IWLSKernel),
+            },
+            name="mu",
+        )
+
+        sigma = lsl.Var.new_param(
+            1.0,
+            lsl.Dist(tfd.InverseGamma, concentration=2.0, scale=1.0),
+            inference=gs.MCMCSpec(gs.IWLSKernel),
+            name="sigma",
+        )
+
+        model = lsl.Model([mu, sigma])
+
+        with pytest.raises(ValueError):
+            gs.LieselMCMC(model).get_kernel_list()
+
+        kernels = gs.LieselMCMC(model, which="a").get_kernel_list()
+        assert isinstance(kernels[0], gs.IWLSKernel)
+        assert isinstance(kernels[1], gs.NUTSKernel)
+
+        kernels = gs.LieselMCMC(model, which="b").get_kernel_list()
+        assert isinstance(kernels[0], gs.IWLSKernel)
+        assert isinstance(kernels[1], gs.IWLSKernel)
+
+    def test_kernel_group(self):
+        spec = gs.MCMCSpec(gs.NUTSKernel, kernel_group="a")
+
+        mu = lsl.Var.new_param(
+            0.0,
+            lsl.Dist(tfd.Normal, loc=0.0, scale=1.0),
+            inference=spec,
+            name="mu",
+        )
+
+        sigma = lsl.Var.new_param(
+            1.0,
+            lsl.Dist(tfd.InverseGamma, concentration=1.0, scale=0.5),
+            inference=spec,
+            name="sigma",
+        )
+
+        model = lsl.Model([mu, sigma])
+
+        mcmc = gs.LieselMCMC(model)
+        kernels = mcmc.get_kernel_list()
+
+        assert len(kernels) == 1
+        assert kernels[0].position_keys == ("sigma", "mu")
+
+    def test_incoherent_kernel_group(self):
+        mu = lsl.Var.new_param(
+            0.0,
+            lsl.Dist(tfd.Normal, loc=0.0, scale=1.0),
+            inference=gs.MCMCSpec(gs.NUTSKernel, kernel_group="a"),
+            name="mu",
+        )
+
+        sigma = lsl.Var.new_param(
+            1.0,
+            lsl.Dist(tfd.InverseGamma, concentration=1.0, scale=0.5),
+            inference=gs.MCMCSpec(gs.IWLSKernel, kernel_group="a"),
+            name="sigma",
+        )
+
+        model = lsl.Model([mu, sigma])
+
+        mcmc = gs.LieselMCMC(model)
+        with pytest.raises(ValueError):
+            mcmc.get_kernel_list()
+
+        mu = lsl.Var.new_param(
+            0.0,
+            lsl.Dist(tfd.Normal, loc=0.0, scale=1.0),
+            inference=gs.MCMCSpec(
+                gs.NUTSKernel, {"da_target_accept": 0.6}, kernel_group="a"
+            ),
+            name="mu",
+        )
+
+        sigma = lsl.Var.new_param(
+            1.0,
+            lsl.Dist(tfd.InverseGamma, concentration=1.0, scale=0.5),
+            inference=gs.MCMCSpec(
+                gs.NUTSKernel, {"da_target_accept": 0.7}, kernel_group="a"
+            ),
+            name="sigma",
+        )
+
+        model = lsl.Model([mu, sigma])
+
+        mcmc = gs.LieselMCMC(model)
+        with pytest.raises(ValueError):
+            mcmc.get_kernel_list()
+
+    def test_jitter_functions(self):
+        mu = lsl.Var.new_param(
+            0.0,
+            lsl.Dist(tfd.Normal, loc=0.0, scale=1.0),
+            inference=gs.MCMCSpec(
+                gs.NUTSKernel,
+                kernel_group="a",
+                jitter_dist=tfd.Uniform(low=-1.0, high=1.0),
+            ),
+            name="mu",
+        )
+
+        sigma = lsl.Var.new_param(
+            1.0,
+            lsl.Dist(tfd.InverseGamma, concentration=1.0, scale=0.5),
+            inference=gs.MCMCSpec(
+                gs.NUTSKernel,
+                kernel_group="a",
+                jitter_dist=tfd.Uniform(low=0.0, high=1.0),
+            ),
+            name="sigma",
+        )
+
+        model = lsl.Model([mu, sigma])
+
+        mcmc = gs.LieselMCMC(model)
+        jitter_fns = mcmc.get_jitter_functions()
+
+        assert len(jitter_fns) == 2
+
+        eb = mcmc.get_engine_builder(1, 4)
+        assert len(eb.jitter_fns.expect("")) == 2
+
+    def test_transform_var_with_inference_new(self):
+        """
+        It is allowed to pass a new inferece object during transformation.
+        In this case, the inference object of the original variable is removed.
+        """
+        sigma = lsl.Var.new_param(
+            1.0,
+            lsl.Dist(tfd.InverseGamma, concentration=2.0, scale=1.0),
+            inference=gs.MCMCSpec(gs.IWLSKernel),
+            name="sigma",
+        )
+        inference = sigma.inference
+
+        log_sigma = sigma.transform(tfb.Exp(), inference=gs.MCMCSpec(gs.NUTSKernel))
+        assert log_sigma.inference is not inference
+        assert sigma.inference is None
+        assert log_sigma.inference.kernel is gs.NUTSKernel
+
+    def test_transform_var_with_inference_none(self):
+        """
+        Default behavior when trying to transform a variable *with* inference
+        information: Error. You need to declare explicitly, what you want to do.
+        In this case, ``"drop"`` means the inference information is deleted.
+        """
+        sigma = lsl.Var.new_param(
+            1.0,
+            lsl.Dist(tfd.InverseGamma, concentration=2.0, scale=1.0),
+            inference=gs.MCMCSpec(gs.IWLSKernel),
+            name="sigma",
+        )
+        inference = sigma.inference
+
+        with pytest.raises(ValueError):
+            sigma.transform(tfb.Exp())
+
+        log_sigma = sigma.transform(tfb.Exp(), inference="drop")
+        assert log_sigma.inference is not inference
+        assert sigma.inference is None
+        assert log_sigma.inference is None
+
+    def test_transform_var_without_inference(self):
+        """
+        Default when the original variable has no inference information:
+        Everything works smoothly.
+        """
+        sigma = lsl.Var.new_param(
+            1.0,
+            lsl.Dist(tfd.InverseGamma, concentration=2.0, scale=1.0),
+            name="sigma",
+        )
+
+        log_sigma = sigma.transform(tfb.Exp())
+
+        assert log_sigma.inference is None
+        assert sigma.inference is None
+        assert log_sigma.inference is None


Reviews:
user: login: copilot-pull-request-reviewer[bot]
body: ## Pull Request Overview

This PR introduces a more flexible and future‐proof approach for configuring MCMC inference within the goose module by adding an inference information container (MCMCSpec) that can support multiple kernel specifications, kernel grouping, and automatic jitter function handling.  
- Introduces the MCMCSpec data class along with multiple inference modes in variables.  
- Updates the Var transformation API to support explicit inference handling and integrates additional tests for these changes.

### Reviewed Changes

Copilot reviewed 4 out of 4 changed files in this pull request and generated 1 comment.

| File                                | Description                                               |
| ----------------------------------- | --------------------------------------------------------- |
| tests/goose/test_builder.py         | New tests for engine construction, multiple specs, and transform behavior. |
| liesel/model/nodes.py               | Updated Var API to support an additional inference attribute and transformations. |
| liesel/goose/builder.py             | Added LieselMCMC class to build and manage MCMC inference workflows. |
| liesel/goose/__init__.py            | Exported the new MCMCSpec and LieselMCMC classes.         |


<details>
<summary>Comments suppressed due to low confidence (1)</summary>

**liesel/goose/builder.py:1708**
* Insert a space to change 'argumentsare' to 'arguments are constant' in the bijector error message.
```
" argumentsare constant."
```
</details>


state: COMMENTED
author_association: CONTRIBUTOR
submitted_at: 2025-04-17T07:43:59Z
user: login: GianmarcoCallegher
state: COMMENTED
author_association: CONTRIBUTOR
submitted_at: 2025-04-23T13:03:27Z
user: login: GianmarcoCallegher
state: COMMENTED
author_association: CONTRIBUTOR
submitted_at: 2025-04-23T13:06:44Z
user: login: GianmarcoCallegher
body: This is great, @jobrachem . I left two small comments
state: COMMENTED
author_association: CONTRIBUTOR
submitted_at: 2025-04-23T13:09:28Z
user: login: jobrachem
state: COMMENTED
author_association: CONTRIBUTOR
submitted_at: 2025-04-23T14:12:30Z
user: login: jobrachem
state: COMMENTED
author_association: CONTRIBUTOR
submitted_at: 2025-04-23T14:12:37Z
user: login: wiep
body: Not much to add after our discussion. I like the idea and implementation. Let's briefly consider to be generic over the inference type.
state: COMMENTED
author_association: CONTRIBUTOR
submitted_at: 2025-04-25T16:47:27Z
user: login: jobrachem
state: COMMENTED
author_association: CONTRIBUTOR
submitted_at: 2025-04-28T14:50:40Z
user: login: jobrachem
state: COMMENTED
author_association: CONTRIBUTOR
submitted_at: 2025-04-28T15:06:24Z
user: login: jobrachem
state: COMMENTED
author_association: CONTRIBUTOR
submitted_at: 2025-04-28T15:15:20Z
user: login: wiep
state: COMMENTED
author_association: CONTRIBUTOR
submitted_at: 2025-04-28T15:39:50Z
user: login: wiep
body: can be merged.
state: APPROVED
author_association: CONTRIBUTOR
submitted_at: 2025-04-28T15:40:56Z

Review comments:
diff_hunk: @@ -1610,6 +1650,12 @@ def transform(original_var: lsl.Var, bijector: tfb.Bijector):
         >>> scale.update().log_prob
         0.0
         """
+        if inference is None and self.inference:
+            raise ValueError(
+                f"{self} has inference information in the .inference attribute. "
+                "To proceed with transformation, the .inference information needs to"
+                "be explicitly removed. You can transform with ``infrence='drop'``."
path: liesel/model/nodes.py
user: login: Copilot
body: Correct the typo 'infrence' to 'inference' in the exception message to ensure clarity.
```suggestion
                "be explicitly removed. You can transform with ``inference='drop'``."
```
created_at: 2025-04-17T07:43:59Z
author_association: CONTRIBUTOR
diff_hunk: @@ -485,3 +495,131 @@ def build(self) -> Engine:
             quantity_generators=self.quantity_generators,
             show_progress=self.show_progress,
         )
+
+
+@dataclass
+class LieselMCMC:
+    model: Model
+    which: str | None = None
+
+    def get_spec(self, var: Var) -> MCMCSpec | None:
+        inference = var.get_inference(self.which)
+        if inference is None:
+            return inference
+
+        if not isinstance(inference, MCMCSpec):
+            raise ValueError(
+                f"Attribute 'inference' of variable {var} is of type"
+                f" {type(inference)}, but expected type '{MCMCSpec}'."
+            )
+        return inference
+
+    def kernel_groups(self) -> dict[str, _KernelGroup]:
+        vars_ = self.model.vars
+        kernel_groups: dict[str, _KernelGroup] = {}
+
+        # collect kernels
+        for name, var in vars_.items():
+            inference = self.get_spec(var)
+
+            if not inference:
+                continue
+
+            group_name = inference.kernel_group
+
+            if group_name is None:  # if this is a kernel without group
+                kernel_groups[name] = _KernelGroup(
+                    kernel=inference.kernel,
+                    kwargs=inference.kernel_kwargs,
+                    position_keys=[name],
+                )
+
+            elif group_name in kernel_groups:  # check group and append var name
+                group = kernel_groups[group_name]
+                same_kernel = group.kernel is inference.kernel
+                if not same_kernel:
+                    raise ValueError(
+                        "Found incoherent kernel classes for kernel group"
+                        f" {group_name}."
+                    )
+
+                same_kwargs = group.kwargs == inference.kernel_kwargs
+                if not same_kwargs:
+                    raise ValueError(
+                        f"Found incoherent kwargs for kernel group {group_name}."
+                    )
+                group.position_keys.append(name)
+
+            else:  # start a new group
+                kernel_groups[group_name] = _KernelGroup(
+                    kernel=inference.kernel,
+                    kwargs=inference.kernel_kwargs,
+                    position_keys=[name],
+                )
+
+        return kernel_groups
+
+    def kernel_list(self) -> list[Kernel]:
+        kernel_groups = self.kernel_groups()
+        kernel_list = [
+            g.kernel(g.position_keys, **g.kwargs)  # type: ignore
+            for g in kernel_groups.values()
+        ]
path: liesel/goose/builder.py
user: login: GianmarcoCallegher
body: Can it be a property? 

I'd name it `get_kernel_list`, otherwise
created_at: 2025-04-23T13:03:27Z
author_association: CONTRIBUTOR
diff_hunk: @@ -485,3 +495,131 @@ def build(self) -> Engine:
             quantity_generators=self.quantity_generators,
             show_progress=self.show_progress,
         )
+
+
+@dataclass
+class LieselMCMC:
+    model: Model
+    which: str | None = None
+
+    def get_spec(self, var: Var) -> MCMCSpec | None:
+        inference = var.get_inference(self.which)
+        if inference is None:
+            return inference
+
+        if not isinstance(inference, MCMCSpec):
+            raise ValueError(
+                f"Attribute 'inference' of variable {var} is of type"
+                f" {type(inference)}, but expected type '{MCMCSpec}'."
+            )
+        return inference
+
+    def kernel_groups(self) -> dict[str, _KernelGroup]:
+        vars_ = self.model.vars
+        kernel_groups: dict[str, _KernelGroup] = {}
path: liesel/goose/builder.py
user: login: GianmarcoCallegher
body: Also here, I would set it as a property or call it `get_kernel_groups`
created_at: 2025-04-23T13:06:44Z
author_association: CONTRIBUTOR
diff_hunk: @@ -485,3 +495,131 @@ def build(self) -> Engine:
             quantity_generators=self.quantity_generators,
             show_progress=self.show_progress,
         )
+
+
+@dataclass
+class LieselMCMC:
+    model: Model
+    which: str | None = None
+
+    def get_spec(self, var: Var) -> MCMCSpec | None:
+        inference = var.get_inference(self.which)
+        if inference is None:
+            return inference
+
+        if not isinstance(inference, MCMCSpec):
+            raise ValueError(
+                f"Attribute 'inference' of variable {var} is of type"
+                f" {type(inference)}, but expected type '{MCMCSpec}'."
+            )
+        return inference
+
+    def kernel_groups(self) -> dict[str, _KernelGroup]:
+        vars_ = self.model.vars
+        kernel_groups: dict[str, _KernelGroup] = {}
+
+        # collect kernels
+        for name, var in vars_.items():
+            inference = self.get_spec(var)
+
+            if not inference:
+                continue
+
+            group_name = inference.kernel_group
+
+            if group_name is None:  # if this is a kernel without group
+                kernel_groups[name] = _KernelGroup(
+                    kernel=inference.kernel,
+                    kwargs=inference.kernel_kwargs,
+                    position_keys=[name],
+                )
+
+            elif group_name in kernel_groups:  # check group and append var name
+                group = kernel_groups[group_name]
+                same_kernel = group.kernel is inference.kernel
+                if not same_kernel:
+                    raise ValueError(
+                        "Found incoherent kernel classes for kernel group"
+                        f" {group_name}."
+                    )
+
+                same_kwargs = group.kwargs == inference.kernel_kwargs
+                if not same_kwargs:
+                    raise ValueError(
+                        f"Found incoherent kwargs for kernel group {group_name}."
+                    )
+                group.position_keys.append(name)
+
+            else:  # start a new group
+                kernel_groups[group_name] = _KernelGroup(
+                    kernel=inference.kernel,
+                    kwargs=inference.kernel_kwargs,
+                    position_keys=[name],
+                )
+
+        return kernel_groups
+
+    def kernel_list(self) -> list[Kernel]:
+        kernel_groups = self.kernel_groups()
+        kernel_list = [
+            g.kernel(g.position_keys, **g.kwargs)  # type: ignore
+            for g in kernel_groups.values()
+        ]
path: liesel/goose/builder.py
user: login: jobrachem
body: Fair point, thanks @GianmarcoCallegher ! I updated all method names on the `gs.LieselMCMC` object. I prefer `get_` methods here to avoid users to think that these are persistent lists that they could append stuff to.

![image](https://github.com/user-attachments/assets/42a5590b-33ff-430f-8024-1b8ce7fddbe1)

created_at: 2025-04-23T14:12:29Z
author_association: CONTRIBUTOR
diff_hunk: @@ -485,3 +495,131 @@ def build(self) -> Engine:
             quantity_generators=self.quantity_generators,
             show_progress=self.show_progress,
         )
+
+
+@dataclass
+class LieselMCMC:
+    model: Model
+    which: str | None = None
+
+    def get_spec(self, var: Var) -> MCMCSpec | None:
+        inference = var.get_inference(self.which)
+        if inference is None:
+            return inference
+
+        if not isinstance(inference, MCMCSpec):
+            raise ValueError(
+                f"Attribute 'inference' of variable {var} is of type"
+                f" {type(inference)}, but expected type '{MCMCSpec}'."
+            )
+        return inference
+
+    def kernel_groups(self) -> dict[str, _KernelGroup]:
+        vars_ = self.model.vars
+        kernel_groups: dict[str, _KernelGroup] = {}
path: liesel/goose/builder.py
user: login: jobrachem
body: Fair point, thanks @GianmarcoCallegher ! I updated all method names on the `gs.LieselMCMC` object. I prefer `get_` methods here to avoid users to think that these are persistent lists that they could append stuff to.

![image](https://github.com/user-attachments/assets/42a5590b-33ff-430f-8024-1b8ce7fddbe1)

created_at: 2025-04-23T14:12:37Z
author_association: CONTRIBUTOR
diff_hunk: @@ -7,21 +7,27 @@
 implementations.
 """
 
+from __future__ import annotations
path: liesel/goose/builder.py
user: login: wiep
body: Postponed evaluation of annotations is default since 3.10 (imho). So, we probably don't need it.
created_at: 2025-04-25T16:29:55Z
author_association: CONTRIBUTOR
diff_hunk: @@ -485,3 +495,131 @@ def build(self) -> Engine:
             quantity_generators=self.quantity_generators,
             show_progress=self.show_progress,
         )
+
+
+@dataclass
+class LieselMCMC:
path: liesel/goose/builder.py
user: login: wiep
body: just wondering whether this code should go in a new module to keep files a bit smaller. what do you think?
created_at: 2025-04-25T16:31:49Z
author_association: CONTRIBUTOR
diff_hunk: @@ -32,8 +32,11 @@
 from .viz import plot_nodes, plot_vars
 
 if TYPE_CHECKING:
+    from ..goose import MCMCSpec
     from .model import Model
 
+    InferenceTypes = Union[MCMCSpec, None, dict[str, MCMCSpec]]
path: liesel/model/nodes.py
user: login: wiep
body: not allowing `any` here, is a bit restrictive if a user has their own inference implementation and wants to use this field. I don't see this as very important right now. However, we could also consider making the types generic over the inference type. Using PEP 695 (v.3.12+) notation that would look like

```python

type InferenceTypes[T] = T | dict[str, T] | None

class Var[InferenceT]:
    def __init__(self, inference: InferenceTypes[InferenceT] = None):
        self->inference = inference
```

Pro: future proof
Neg: more complicated / support of generics is not as nice in < 3.12.

created_at: 2025-04-25T16:39:52Z
author_association: CONTRIBUTOR
diff_hunk: @@ -7,21 +7,27 @@
 implementations.
 """
 
+from __future__ import annotations
path: liesel/goose/builder.py
user: login: jobrachem
body: flake8 fails in the pre-commit if I remove the line, so for now I would keep it.
created_at: 2025-04-28T14:50:40Z
author_association: CONTRIBUTOR
diff_hunk: @@ -485,3 +495,131 @@ def build(self) -> Engine:
             quantity_generators=self.quantity_generators,
             show_progress=self.show_progress,
         )
+
+
+@dataclass
+class LieselMCMC:
path: liesel/goose/builder.py
user: login: jobrachem
body: Sure, I moved the new goose code to `goose/mcmc_spec.py`
created_at: 2025-04-28T15:06:23Z
author_association: CONTRIBUTOR
diff_hunk: @@ -32,8 +32,11 @@
 from .viz import plot_nodes, plot_vars
 
 if TYPE_CHECKING:
+    from ..goose import MCMCSpec
     from .model import Model
 
+    InferenceTypes = Union[MCMCSpec, None, dict[str, MCMCSpec]]
path: liesel/model/nodes.py
user: login: jobrachem
body: The PEP 695 (v.3.12+) notation is pretty cool! I did not know about that before. I would like to have that at some point. For now, I'd say let's go forward as it is, but open an issue, ok? I added `Any` to the `InferenceTypes`.
created_at: 2025-04-28T15:15:20Z
author_association: CONTRIBUTOR
diff_hunk: @@ -7,21 +7,27 @@
 implementations.
 """
 
+from __future__ import annotations
path: liesel/goose/builder.py
user: login: wiep
body: Apparently, stringyfying annotations is not the default ([see 3.11 release notes](https://docs.python.org/3/whatsnew/3.11.html#pep-563-may-not-be-the-future)). Instead they are going for a [pep 649](https://peps.python.org/pep-0649/) to be default in 3.14.

I guess that means, we can either use the import or manually stringyfy annotations.
created_at: 2025-04-28T15:39:50Z
author_association: CONTRIBUTOR

PR comments:
user: login: jobrachem
created_at: 2025-04-05T07:02:29Z
body: ## Behavior during transformation

```python
sigma = lsl.Var.new_param(
    1.0,
    lsl.Dist(tfd.InverseGamma, concentration=2.0, scale=1.0),
    inference=gs.MCMCSpec(gs.IWLSKernel),
    name="sigma",
)

log_sigma = sigma.transform(tfb.Exp(), inference=None) # default: No inference info
log_sigma = sigma.transform(tfb.Exp(), inference=gs.MCMCSpec(gs.NUTSKernel)) # manual
log_sigma = sigma.transform(tfb.Exp(), inference="transfer") # transfer from sigma

sigma.inference # is always None after transformation
```
author_association: CONTRIBUTOR
user: login: jobrachem
created_at: 2025-04-05T12:57:06Z
body: @GianmarcoCallegher @wiep 
I drafted a 'maximalist' implementation of the more abstract functionality that we talked about yesterday here. Turns out, I really like this way of setting things up. Thanks for the discussion yesterday! Please let me know your thoughts about this draft when you find the time. Otherwise, of course we can continue the discussion in the next weekly.
author_association: CONTRIBUTOR
user: login: jobrachem
created_at: 2025-04-15T13:53:53Z
body: - `Var.transform(inference="drop")` for dropping explicitly; `Var.transform(inference=None)` raises error if there is `Var.inference`.

- Check what happens when `inference="transfer"`, but `Var.inference is None`. Should be an error. 

```
a = Var(...)

inference = a.inference
a.inference = None
b = a.transform()
    a.inference.allow_transformation
b.inference = inference
```


```
a.transform(inference="drop")
```
author_association: CONTRIBUTOR
user: login: wiep
created_at: 2025-04-15T17:13:37Z
body: Thoughts from the meeting. Transforming the inference part of `a.transform` could be handled by the object in `a.inference` as this object may have access to additional information about the inference scheme.




author_association: CONTRIBUTOR
user: login: jobrachem
created_at: 2025-04-17T07:42:33Z
body: Updated behavior as discussed; described by the three tests below. @wiep @GianmarcoCallegher this is ready for review now :)

```python
def test_transform_var_with_inference_new(self):
    """
    It is allowed to pass a new inferece object during transformation.
    In this case, the inference object of the original variable is removed.
    """
    sigma = lsl.Var.new_param(
        1.0,
        lsl.Dist(tfd.InverseGamma, concentration=2.0, scale=1.0),
        inference=gs.MCMCSpec(gs.IWLSKernel),
        name="sigma",
    )
    inference = sigma.inference

    log_sigma = sigma.transform(tfb.Exp(), inference=gs.MCMCSpec(gs.NUTSKernel))
    assert log_sigma.inference is not inference
    assert sigma.inference is None
    assert log_sigma.inference.kernel is gs.NUTSKernel

def test_transform_var_with_inference_none(self):
    """
    Default behavior when trying to transform a variable *with* inference 
    information: Error. You need to declare explicitly, what you want to do.
    In this case, ``"drop"`` means the inference information is deleted.
    """
    sigma = lsl.Var.new_param(
        1.0,
        lsl.Dist(tfd.InverseGamma, concentration=2.0, scale=1.0),
        inference=gs.MCMCSpec(gs.IWLSKernel),
        name="sigma",
    )
    inference = sigma.inference

    with pytest.raises(ValueError):
        sigma.transform(tfb.Exp())

    log_sigma = sigma.transform(tfb.Exp(), inference="drop")
    assert log_sigma.inference is not inference
    assert sigma.inference is None
    assert log_sigma.inference is None

def test_transform_var_without_inference(self):
    """
    Default when the original variable has no inference information: 
    Everything works smoothly.
    """
    sigma = lsl.Var.new_param(
        1.0,
        lsl.Dist(tfd.InverseGamma, concentration=2.0, scale=1.0),
        name="sigma",
    )
    
    log_sigma = sigma.transform(tfb.Exp())
    
    assert log_sigma.inference is None
    assert sigma.inference is None
    assert log_sigma.inference is None
```
author_association: CONTRIBUTOR
user: login: jobrachem
created_at: 2025-04-23T09:49:20Z
body: @wiep @GianmarcoCallegher I updated the PR by removing the duration-arguments from the `gs.LieselMCMC.engine_builder()` method. This way, the `.engine_builder()` method does not include an assumption about the structure of the warmup scheme, and users are nudged to set it themselves with `EngineBuilder.set_duration()` or `EngineBuilder.set_epochs()`.

For the same reason I removed `gs.LieselMCMC.engine()`, because returning a built engine requires assumptions about the structure of the warmup.
author_association: CONTRIBUTOR