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

Author: zaxtax

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,25 @@ 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 in the original
+        (constrained) space of the model.
+    std : Dataset or None
+        Posterior standard deviation of each latent variable in the
+        original (constrained) space. ``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
@@ -1172,6 +1192,79 @@ def std_data(self) -> Dataset:
         """Standard deviation of the latent variables as an xarray Dataset."""
         return self.var_to_data(self.std)
 
+    def _untransform_tensor(self, flat_array: np.ndarray) -> Dataset:
+        """Map a flat array in unconstrained space to an xarray Dataset in the original (constrained) space via the model transforms.
+
+        Parameters
+        ----------
+        flat_array : np.ndarray
+            Flat parameter vector in unconstrained (transformed) space.
+
+        Returns
+        -------
+        Dataset
+            Named variables in the original constrained space.
+        """
+        result = {}
+        for rv in self.group:
+            transform = self.model.rvs_to_transforms.get(rv)
+            value_var = self.model.rvs_to_values[rv]
+            name = value_var.name
+            _, slc, shape, dtype = self.ordering[name]
+
+            unconstrained = flat_array[slc].reshape(shape).astype(dtype)
+
+            if transform is not None:
+                constrained = transform.backward(
+                    pt.as_tensor(unconstrained), *rv.owner.inputs
+                ).eval()
+            else:
+                constrained = unconstrained
+
+            rv_name = rv.name
+            dims = self.model.named_vars_to_dims.get(rv_name, None)
+            if dims is not None:
+                coords = {d: np.array(self.model.coords[d]) for d in dims}
+            else:
+                coords = None
+            result[rv_name] = DataArray(constrained, coords=coords, dims=dims, name=rv_name)
+
+        return Dataset(result)
+
+    @property
+    def untransformed_mean_data(self) -> Dataset:
+        """Mean of the latent variables in the original (untransformed) space.
+
+        Returns
+        -------
+        Dataset
+            xarray Dataset with variable names from the original model space.
+        """
+        return self._untransform_tensor(self.mean.eval())
+
+    @property
+    def untransformed_std_data(self) -> Dataset:
+        """Standard deviation of the latent variables in the original (untransformed) space.
+
+        Note that the standard deviation in the original space is computed
+        from the unconstrained parameters and may not directly correspond
+        to the standard deviation of the constrained posterior.
+
+        Returns
+        -------
+        Dataset
+            xarray Dataset with variable names from the original model space.
+        """
+        return self._untransform_tensor(self.std.eval())
+
+    @property
+    def state(self) -> VIState:
+        """Fit state with mean and std in the original (constrained) space."""
+        return VIState(
+            mean=self.untransformed_mean_data,
+            std=self.untransformed_std_data if self.has_logq else None,
+        )
+
 
 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
 
@@ -503,3 +505,131 @@ def test_multiple_minibatch_variables():
         )
         mean_field = pm.fit(10_000, obj_optimizer=pm.adam(learning_rate=0.01), progressbar=False)
     np.testing.assert_allclose(mean_field.mean.get_value(), true_weights, rtol=1e-1)
+
+
+class TestUntransformedData:
+    def test_untransformed_keys(self, hierarchical_model, hierarchical_model_data):
+        """Test that untransformed_mean_data uses original RV names."""
+        with hierarchical_model:
+            fitted = pm.fit(1, progressbar=False)
+
+        assert "sigma_log__" in fitted.mean_data
+        ut = fitted.untransformed_mean_data
+        assert set(ut.keys()) == {"sigma", "sigma_group_mu", "group_mu", "mu"}
+        assert ut["group_mu"].shape == hierarchical_model_data["group_shape"]
+        assert list(ut["group_mu"].coords.keys()) == list(
+            hierarchical_model_data["group_coords"].keys()
+        )
+
+    def test_log_transform_inversion(self):
+        """HalfNormal: verify backward via exp."""
+        rng = np.random.default_rng(42)
+        with pm.Model():
+            sigma = pm.HalfNormal("sigma", sigma=5.0)
+            mu = pm.Normal("mu", 0, 1)
+            pm.Normal("y", mu, sigma, observed=rng.normal(size=3))
+            fitted = pm.fit(100, progressbar=False, random_seed=42)
+
+        sigma_log_mean = fitted.mean_data["sigma_log__"].values
+        sigma_mean = fitted.untransformed_mean_data["sigma"].values
+        np.testing.assert_allclose(sigma_mean, np.exp(sigma_log_mean), rtol=1e-6)
+
+    def test_no_transform_passthrough(self, simple_model):
+        """Variables without transforms have identical values in both spaces."""
+        with simple_model:
+            fitted = pm.fit(100, progressbar=False, random_seed=42)
+
+        transformed_mu = fitted.mean_data["mu"].values
+        untransformed_mu = fitted.untransformed_mean_data["mu"].values
+        np.testing.assert_array_equal(transformed_mu, untransformed_mu)
+
+    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"