Adds a state property that mean and std are guaranteed to be on. (#8319)

* Adds a state property that mean and std are guaranteed to be on.

Uses the recently merged constrain_values utility from
pymc.model.transform_values to map unconstrained parameters
back to the original constrained space, replacing the custom
_untransform_tensor helper.

Author: zaxtax

pymc/variational/inference.py

@@ -112,6 +112,7 @@ def fit(
         progressbar_theme=default_progress_theme,
         *,
         backend=None,
+        include_transformed=False,
         **kwargs,
     ):
         """Perform Operator Variational Inference.
@@ -180,6 +181,9 @@ def fit(
         self.approx.hist = self.hist
         self.state = state
 
+        for group in self.approx.groups:
+            group.include_transformed = include_transformed
+
         return self.approx
 
     def _iterate_without_loss(self, s, n, step_func, progressbar, progressbar_theme, callbacks):
@@ -698,6 +702,7 @@ def fit(
     inf_kwargs=None,
     *,
     backend=None,
+    include_transformed=False,
     **kwargs,
 ):
     r"""Handy shortcut for using inference methods in functional way.
@@ -725,6 +730,10 @@ def fit(
         starting standard deviation for inference, only available for method 'advi'
     backend: str, optional
         Which computational backend to use. Recommended to be one of "numba", "c", and "jax".
+    include_transformed: bool, default False
+        If True, the :meth:`~Approximation.state` will also include the
+        unconstrained (transformed) variables alongside the original model
+        variables, similar to ``pm.sample(idata_kwargs={'include_transformed': True})``.
 
     Other Parameters
     ----------------
@@ -787,4 +796,4 @@ def fit(
         inference = method
     else:
         raise TypeError(f"method should be one of {set(_select.keys())} or Inference instance")
-    return inference.fit(n, backend=backend, **kwargs)
+    return inference.fit(n, backend=backend, include_transformed=include_transformed, **kwargs)

pymc/variational/opvi.py

@@ -52,6 +52,7 @@
 import itertools
 import warnings
 
+from dataclasses import dataclass
 from typing import Any, overload
 
 import numpy as np
@@ -117,6 +118,23 @@ class GroupError(VariationalInferenceError, TypeError):
     """Error related to VI groups."""
 
 
+@dataclass
+class VIState:
+    """State of a fitted variational inference approximation.
+
+    Parameters
+    ----------
+    mean : Dataset
+        Posterior mean of each latent variable.
+    std : Dataset or None
+        Posterior standard deviation of each latent variable.
+        ``None`` for particle-based methods.
+    """
+
+    mean: Dataset
+    std: Dataset | None
+
+
 def _known_scan_ignored_inputs(terms):
     # TODO: remove when scan issue with grads is fixed
     from pymc.data import MinibatchOp
@@ -1164,14 +1182,53 @@ def var_to_data(self, shared: pt.TensorVariable) -> Dataset:
 
     @property
     def mean_data(self) -> Dataset:
-        """Mean of the latent variables as an xarray Dataset."""
+        """Mean of the latent variables as an xarray Dataset.
+
+        These are the values the optimizer works with. For the
+        original model space, use :meth:`state` instead.
+        """
         return self.var_to_data(self.mean)
 
     @property
     def std_data(self) -> Dataset:
-        """Standard deviation of the latent variables as an xarray Dataset."""
+        """Standard deviation of the latent variables as an xarray Dataset.
+
+        These are the values the optimizer works with. For the
+        original model space, use :meth:`state` instead.
+        """
         return self.var_to_data(self.std)
 
+    @property
+    def state(self) -> VIState:
+        """Fit state with mean and std as xarray Datasets."""
+        from pymc.model.transform_values import constrain_values
+
+        include_transformed = getattr(self, "include_transformed", False)
+
+        def _constrain_flat(flat_tensor: pt.TensorVariable) -> Dataset:
+            ds = self.var_to_data(flat_tensor)
+            ds = ds.expand_dims("__sample__")
+            result = constrain_values(
+                ds,
+                model=self.model,
+                sample_dims=("__sample__",),
+                compile_kwargs={"mode": "FAST_COMPILE"},
+            )
+            return result.squeeze("__sample__", drop=True)
+
+        mean_ds = _constrain_flat(self.mean)
+        std_ds = _constrain_flat(self.std) if self.has_logq else None
+
+        if include_transformed:
+            mean_ds = mean_ds.merge(self.mean_data, compat="override")
+            if std_ds is not None:
+                std_ds = std_ds.merge(self.std_data, compat="override")
+
+        return VIState(
+            mean=mean_ds,
+            std=std_ds,
+        )
+
 
 group_for_params = Group.group_for_params
 group_for_short_name = Group.group_for_short_name

tests/variational/test_inference.py

@@ -22,6 +22,8 @@
 import pytensor.tensor as pt
 import pytest
 
+from xarray import Dataset
+
 import pymc as pm
 import pymc.variational.opvi as opvi
 
@@ -509,9 +511,163 @@ def test_sample_outside_model_context():
     with pm.Model() as model:
         mu = pm.Normal("mu", 0, 1)
 
-    # Fit with explicit model, then exit the model context
     approx = pm.fit(10, method="advi", model=model, progressbar=False)
-
-    # sample() should work without an active model context
     trace = approx.sample(50)
     assert trace.posterior["mu"].shape == (1, 50)
+
+
+class TestUntransformedData:
+    def test_state_mean_field(self):
+        """ADVI state has family='mean_field', mean and std in constrained space."""
+        rng = np.random.default_rng(42)
+        with pm.Model():
+            pm.HalfNormal("sigma", sigma=5.0)
+            pm.Normal("mu", 0, 1)
+            pm.Normal("y", rng.normal(size=3), observed=rng.normal(size=3))
+            fitted = pm.fit(100, method="advi", progressbar=False, random_seed=42)
+
+        s = fitted.state
+        assert set(s.mean.keys()) == {"sigma", "mu"}
+        assert set(s.std.keys()) == {"sigma", "mu"}
+        assert s.std is not None
+        assert s.mean["sigma"].values > 0
+        assert s.std["sigma"].values > 0
+
+    def test_state_full_rank(self):
+        """FullRankADVI state has mean and std."""
+        rng = np.random.default_rng(42)
+        with pm.Model():
+            pm.HalfNormal("sigma", sigma=5.0)
+            pm.Normal("mu", 0, 1)
+            pm.Normal("y", rng.normal(size=3), observed=rng.normal(size=3))
+            fitted = pm.fit(100, method="fullrank_advi", progressbar=False, random_seed=42)
+
+        s = fitted.state
+        assert s.mean.keys() == {"sigma", "mu"}
+        assert s.std is not None
+        assert s.mean["sigma"].values > 0
+
+    def test_state_empirical_std_is_none(self):
+        """Empirical state has std=None."""
+        rng = np.random.default_rng(42)
+        with pm.Model():
+            pm.Normal("mu", 0, 1)
+            pm.Normal("y", rng.normal(size=10), observed=rng.normal(size=10))
+            inference = pm.SVGD(n_particles=50, random_seed=42)
+            fitted = inference.fit(100, progressbar=False)
+
+        s = fitted.state
+        assert s.std is None
+        assert "mu" in s.mean
+
+    def test_state_is_single_group_approx_attr(self):
+        """state is accessible from SingleGroupApproximation via __getattr__ proxy."""
+        with pm.Model():
+            pm.Normal("mu", 0, 1)
+            inference = pm.ADVI(random_seed=42)
+            fitted = inference.fit(10, progressbar=False)
+
+        s = fitted.state
+        assert "mu" in s.mean
+
+    def test_state_in_callback(self):
+        """Callbacks can access state during training."""
+        rng = np.random.default_rng(42)
+        snapshots = []
+
+        def callback(approx, losses, i):
+            s = approx.state
+            snapshots.append(
+                {
+                    "i": i,
+                    "mean": s.mean,
+                    "std": s.std,
+                }
+            )
+
+        with pm.Model():
+            pm.HalfNormal("sigma", sigma=5.0)
+            pm.Normal("mu", 0, 1)
+            pm.Normal("y", rng.normal(size=3), observed=rng.normal(size=3))
+            inference = pm.ADVI(random_seed=42)
+            fitted = inference.fit(50, callbacks=[callback], progressbar=False)
+
+        assert len(snapshots) == 50
+        for snap in snapshots:
+            assert isinstance(snap["mean"], Dataset)
+            assert set(snap["mean"].keys()) == {"sigma", "mu"}
+            assert snap["std"] is not None
+            assert set(snap["std"].keys()) == {"sigma", "mu"}
+        # The last snapshot should match the final state
+        final = fitted.state
+        np.testing.assert_allclose(
+            snapshots[-1]["mean"]["sigma"].values, final.mean["sigma"].values
+        )
+        # Parameters should have moved from their initial values
+        first_mean = snapshots[0]["mean"]["mu"].values
+        last_mean = snapshots[-1]["mean"]["mu"].values
+        assert not np.allclose(first_mean, last_mean), "parameters should change during training"
+
+    def test_state_dirichlet(self):
+        """State works with Dirichlet (simplex transform changes dimensionality)."""
+        with pm.Model():
+            pm.Dirichlet("p", a=[1, 2, 3])
+            fitted = pm.fit(50, method="advi", progressbar=False, random_seed=42)
+
+        s = fitted.state
+        # Dirichlet with K=3 has K-1=2 unconstrained dims, 3 constrained dims
+        assert "p" in s.mean
+        assert s.mean["p"].values.shape == (3,)
+        # Values should be on the simplex (sum to 1)
+        np.testing.assert_allclose(s.mean["p"].values.sum(), 1.0, atol=1e-6)
+        assert (s.mean["p"].values >= 0).all()
+        assert (s.mean["p"].values <= 1).all()
+        # std should also be in constrained space
+        assert s.std is not None
+        assert "p" in s.std
+        assert s.std["p"].values.shape == (3,)
+
+    def test_state_include_transformed(self):
+        """include_transformed=True adds unconstrained variables to state."""
+        with pm.Model():
+            pm.HalfNormal("sigma", sigma=5.0)
+            pm.Normal("mu", 0, 1)
+            fitted = pm.fit(
+                50,
+                method="advi",
+                progressbar=False,
+                random_seed=42,
+                include_transformed=True,
+            )
+
+        s = fitted.state
+        # Constrained variables always present
+        assert "sigma" in s.mean
+        assert "mu" in s.mean
+        # Unconstrained variables included when include_transformed=True
+        assert "sigma_log__" in s.mean
+        # mu has no transform, so it appears only once
+        assert list(s.mean.data_vars) == ["sigma", "mu", "sigma_log__"]
+        assert s.std is not None
+        assert "sigma_log__" in s.std
+
+    def test_state_include_transformed_dirichlet(self):
+        """include_transformed=True with Dirichlet (dimensionality-changing transform)."""
+        with pm.Model():
+            pm.Dirichlet("p", a=[1, 2, 3])
+            fitted = pm.fit(
+                50,
+                method="advi",
+                progressbar=False,
+                random_seed=42,
+                include_transformed=True,
+            )
+
+        s = fitted.state
+        # Constrained: p (shape 3, on simplex)
+        assert "p" in s.mean
+        assert s.mean["p"].values.shape == (3,)
+        np.testing.assert_allclose(s.mean["p"].values.sum(), 1.0, atol=1e-6)
+        # Unconstrained: p_simplex__ (shape 2, K-1 dims)
+        assert "p_simplex__" in s.mean
+        assert s.mean["p_simplex__"].values.shape == (2,)