add an exceptionn to prevent using SVGD with minibatch

Author: Maxim Kochurov

pymc/tests/variational/test_inference.py

@@ -153,7 +153,10 @@ def fit_kwargs(inference, use_minibatch):
 
 
 def test_fit_oo(inference, fit_kwargs, simple_model_data):
-    trace = inference.fit(**fit_kwargs).sample(10000)
+    try:
+        trace = inference.fit(**fit_kwargs).sample(10000)
+    except opvi.NotImplementedInference as e:
+        pytest.skip(str(e))
     mu_post = simple_model_data["mu_post"]
     d = simple_model_data["d"]
     np.testing.assert_allclose(np.mean(trace.posterior["mu"]), mu_post, rtol=0.05)

pymc/variational/approximations.py

@@ -388,7 +388,5 @@ def sample(post, *_):
         nodes, _ = aesara.scan(
             sample,
             self.histogram,
-            non_sequences=opvi._known_scan_ignored_inputs(node),
-            strict=True,
         )
         return nodes

pymc/variational/opvi.py

@@ -169,8 +169,8 @@ def _known_scan_ignored_inputs(terms):
 
     return [
         n.owner.inputs[0]
-        for n in find_rng_nodes(terms)
-        if isinstance(n, (MinibatchIndexRV, SimulatorRV))
+        for n in aesara.graph.ancestors(terms)
+        if n.owner is not None and isinstance(n.owner.op, (MinibatchIndexRV, SimulatorRV))
     ]
 
 
@@ -1025,9 +1025,7 @@ def symbolic_sample_over_posterior(self, node):
         def sample(post, *_):
             return aesara.clone_replace(node, {self.input: post})
 
-        nodes, _ = aesara.scan(
-            sample, random, non_sequences=_known_scan_ignored_inputs(node), strict=True
-        )
+        nodes, _ = aesara.scan(sample, random)
         assert self.input not in set(aesara.graph.graph_inputs(as_list(nodes)))
         return nodes
 
@@ -1388,8 +1386,6 @@ def sample(*post):
         nodes, _ = aesara.scan(
             sample,
             self.symbolic_randoms,
-            non_sequences=_known_scan_ignored_inputs(node),
-            strict=True,
         )
         assert not (set(self.inputs) & set(aesara.graph.graph_inputs(as_list(nodes))))
         return nodes

pymc/variational/stein.py

@@ -17,7 +17,11 @@
 
 from pymc.aesaraf import floatX
 from pymc.util import WithMemoization, locally_cachedmethod
-from pymc.variational.opvi import node_property
+from pymc.variational.opvi import (
+    NotImplementedInference,
+    _known_scan_ignored_inputs,
+    node_property,
+)
 from pymc.variational.test_functions import rbf
 
 __all__ = ["Stein"]
@@ -46,7 +50,12 @@ def approx_symbolic_matrices(self):
 
     @node_property
     def dlogp(self):
-        grad = at.grad(self.logp_norm.sum(), self.approx_symbolic_matrices)
+        logp = self.logp_norm.sum()
+        if _known_scan_ignored_inputs([logp]):
+            raise NotImplementedInference(
+                "SVGD does not currently support Minibatch or Simulator RV"
+            )
+        grad = at.grad(logp, self.approx_symbolic_matrices)
 
         def flatten2(tensor):
             return tensor.flatten(2)