explicit error about empirical approx (#5874)

* explicit error about empirical approx

* fix an error message

* Write more detailed error message

Co-authored-by: Michael Osthege <[email protected]>

Author: ferrine

RELEASE-NOTES.md

@@ -2,6 +2,7 @@
 
 ## PyMC 4.0.1 (vNext)
 + Fixed an incorrect entry in `pm.Metropolis.stats_dtypes` (see #5582).
++ Added a check in `Empirical` approximation which does not yet support `InferenceData` inputs (see #5874, #5884).
 + ...
 
 ## PyMC 4.0.0 (2022-06-03)

pymc/tests/test_variational_inference.py

@@ -791,6 +791,14 @@ def test_empirical_from_trace(another_simple_model):
         assert emp.histogram.shape[0].eval() == 400
 
 
+def test_empirical_does_not_support_inference_data(another_simple_model):
+    with another_simple_model:
+        step = pm.Metropolis()
+        trace = pm.sample(100, step=step, chains=1, tune=0, return_inferencedata=True)
+        with pytest.raises(NotImplementedError, match="return_inferencedata=False"):
+            Empirical(trace)
+
+
 @pytest.mark.parametrize("score", [True, False])
 def test_fit_with_nans(score):
     X_mean = pm.floatX(np.linspace(0, 10, 10))

pymc/variational/approximations.py

@@ -18,6 +18,7 @@
 from aesara import tensor as at
 from aesara.graph.basic import Variable
 from aesara.tensor.var import TensorVariable
+from arviz import InferenceData
 
 import pymc as pm
 
@@ -198,7 +199,6 @@ def create_shared_params(self, trace=None, size=None, jitter=1, start=None):
                 # Initialize particles
                 histogram = np.tile(start, (size, 1))
                 histogram += pm.floatX(np.random.normal(0, jitter, histogram.shape))
-
         else:
             histogram = np.empty((len(trace) * len(trace.chains), self.ddim))
             i = 0
@@ -210,7 +210,13 @@ def create_shared_params(self, trace=None, size=None, jitter=1, start=None):
 
     def _check_trace(self):
         trace = self._kwargs.get("trace", None)
-        if trace is not None and not all(
+        if isinstance(trace, InferenceData):
+            raise NotImplementedError(
+                "The `Empirical` approximation does not yet support `InferenceData` inputs."
+                " Pass `pm.sample(return_inferencedata=False)` to get a `MultiTrace` to use with `Empirical`."
+                " Please help us to refactor: https://github.com/pymc-devs/pymc/issues/5884"
+            )
+        elif trace is not None and not all(
             [self.model.rvs_to_values[var].name in trace.varnames for var in self.group]
         ):
             raise ValueError("trace has not all free RVs in the group")