add test

Author: ferrine

pymc/variational/approximations.py

@@ -228,7 +228,13 @@ def create_shared_params(self, trace=None, size=None, jitter=1, start=None):
                 for j in range(len(trace)):
                     histogram[i] = DictToArrayBijection.map(trace.point(j, t)).data
                     i += 1
-        return dict(histogram=pytensor.shared(pm.floatX(histogram), "histogram"))
+        return dict(
+            histogram=pytensor.shared(
+                pm.floatX(histogram),
+                "histogram",
+                shape=histogram.shape,
+            )
+        )
 
     def _check_trace(self):
         trace = self._kwargs.get("trace", None)

tests/sampling/test_jax.py

@@ -471,3 +471,10 @@ def test_sample_partially_observed():
     assert idata.observed_data["x_observed"].shape == (2,)
     assert idata.posterior["x_unobserved"].shape == (1, 10, 1)
     assert idata.posterior["x"].shape == (1, 10, 3)
+
+
+@pytest.mark.parametrize("method", ["advi", "fullrank_advi"])
+def test_vi_sampling_jax(method):
+    with pm.Model() as model:
+        x = pm.Normal("x")
+        pm.fit(10, method=method, fn_kwargs=dict(mode="JAX"))