add test

Author: ferrine

pymc/variational/approximations.py

@@ -229,7 +229,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

@@ -84,12 +84,17 @@ def test_jax_PosDefMatrix():
     [
         pytest.param(1),
         pytest.param(
-            2, marks=pytest.mark.skipif(len(jax.devices()) < 2, reason="not enough devices")
+            2,
+            marks=pytest.mark.skipif(
+                len(jax.devices()) < 2, reason="not enough devices"
+            ),
         ),
     ],
 )
 @pytest.mark.parametrize("postprocessing_vectorize", ["scan", "vmap"])
-def test_transform_samples(sampler, postprocessing_backend, chains, postprocessing_vectorize):
+def test_transform_samples(
+    sampler, postprocessing_backend, chains, postprocessing_vectorize
+):
     pytensor.config.on_opt_error = "raise"
     np.random.seed(13244)
 
@@ -236,7 +241,9 @@ def test_replace_shared_variables():
     x = pytensor.shared(5, name="shared_x")
 
     new_x = _replace_shared_variables([x])
-    shared_variables = [var for var in graph_inputs(new_x) if isinstance(var, SharedVariable)]
+    shared_variables = [
+        var for var in graph_inputs(new_x) if isinstance(var, SharedVariable)
+    ]
     assert not shared_variables
 
     x.default_update = x + 1
@@ -263,7 +270,11 @@ def test_get_jaxified_logp():
 @pytest.fixture(scope="module")
 def model_test_idata_kwargs() -> pm.Model:
     with pm.Model(
-        coords={"x_coord": ["a", "b"], "x_coord2": [1, 2], "z_coord": ["apple", "banana", "orange"]}
+        coords={
+            "x_coord": ["a", "b"],
+            "x_coord2": [1, 2],
+            "z_coord": ["apple", "banana", "orange"],
+        }
     ) as m:
         x = pm.Normal("x", shape=(2,), dims=["x_coord"])
         _ = pm.Normal("y", x, observed=[0, 0])
@@ -322,23 +333,30 @@ def test_idata_kwargs(
 
     posterior = idata.get("posterior")
     assert posterior is not None
-    x_dim_expected = idata_kwargs.get("dims", model_test_idata_kwargs.named_vars_to_dims)["x"][0]
+    x_dim_expected = idata_kwargs.get(
+        "dims", model_test_idata_kwargs.named_vars_to_dims
+    )["x"][0]
     assert x_dim_expected is not None
     assert posterior["x"].dims[-1] == x_dim_expected
 
-    x_coords_expected = idata_kwargs.get("coords", model_test_idata_kwargs.coords)[x_dim_expected]
+    x_coords_expected = idata_kwargs.get("coords", model_test_idata_kwargs.coords)[
+        x_dim_expected
+    ]
     assert x_coords_expected is not None
     assert list(x_coords_expected) == list(posterior["x"].coords[x_dim_expected].values)
 
     assert posterior["z"].dims[2] == "z_coord"
     assert np.all(
-        posterior["z"].coords["z_coord"].values == np.array(["apple", "banana", "orange"])
+        posterior["z"].coords["z_coord"].values
+        == np.array(["apple", "banana", "orange"])
     )
 
 
 def test_get_batched_jittered_initial_points():
     with pm.Model() as model:
-        x = pm.MvNormal("x", mu=np.zeros(3), cov=np.eye(3), shape=(2, 3), initval=np.zeros((2, 3)))
+        x = pm.MvNormal(
+            "x", mu=np.zeros(3), cov=np.eye(3), shape=(2, 3), initval=np.zeros((2, 3))
+        )
 
     # No jitter
     ips = _get_batched_jittered_initial_points(
@@ -347,13 +365,17 @@ def test_get_batched_jittered_initial_points():
     assert np.all(ips[0] == 0)
 
     # Single chain
-    ips = _get_batched_jittered_initial_points(model=model, chains=1, random_seed=1, initvals=None)
+    ips = _get_batched_jittered_initial_points(
+        model=model, chains=1, random_seed=1, initvals=None
+    )
 
     assert ips[0].shape == (2, 3)
     assert np.all(ips[0] != 0)
 
     # Multiple chains
-    ips = _get_batched_jittered_initial_points(model=model, chains=2, random_seed=1, initvals=None)
+    ips = _get_batched_jittered_initial_points(
+        model=model, chains=2, random_seed=1, initvals=None
+    )
 
     assert ips[0].shape == (2, 2, 3)
     assert np.all(ips[0][0] != ips[0][1])
@@ -372,7 +394,10 @@ def test_get_batched_jittered_initial_points():
     [
         pytest.param(1),
         pytest.param(
-            2, marks=pytest.mark.skipif(len(jax.devices()) < 2, reason="not enough devices")
+            2,
+            marks=pytest.mark.skipif(
+                len(jax.devices()) < 2, reason="not enough devices"
+            ),
         ),
     ],
 )
@@ -396,8 +421,12 @@ def test_seeding(chains, random_seed, sampler):
         assert all_equal
 
     if chains > 1:
-        assert np.all(result1.posterior["x"].sel(chain=0) != result1.posterior["x"].sel(chain=1))
-        assert np.all(result2.posterior["x"].sel(chain=0) != result2.posterior["x"].sel(chain=1))
+        assert np.all(
+            result1.posterior["x"].sel(chain=0) != result1.posterior["x"].sel(chain=1)
+        )
+        assert np.all(
+            result2.posterior["x"].sel(chain=0) != result2.posterior["x"].sel(chain=1)
+        )
 
 
 @mock.patch("numpyro.infer.MCMC")
@@ -503,3 +532,10 @@ def test_convergence_warnings(caplog, nuts_sampler):
 
     [record] = caplog.records
     assert re.match(r"There were \d+ divergences after tuning", record.message)
+
+
+@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"))