update pytensor version, make xfail more elaborate

Author: ferrine

tests/sampling/test_jax.py

@@ -86,16 +86,12 @@ def test_jax_PosDefMatrix():
         pytest.param(1),
         pytest.param(
             2,
-            marks=pytest.mark.skipif(
-                len(jax.devices()) < 2, reason="not enough devices"
-            ),
+            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)
 
@@ -242,9 +238,7 @@ 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
@@ -332,30 +326,23 @@ 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(
@@ -364,17 +351,13 @@ 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])
@@ -394,9 +377,7 @@ 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"
-            ),
+            marks=pytest.mark.skipif(len(jax.devices()) < 2, reason="not enough devices"),
         ),
     ],
 )
@@ -420,12 +401,8 @@ 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")
@@ -555,7 +532,21 @@ def test_vi_sampling_jax(method):
         pm.fit(10, method=method, fn_kwargs=dict(mode="JAX"))
 
 
-@pytest.mark.xfail(reason="Due to https://github.com/pymc-devs/pytensor/issues/595")
+@pytest.mark.xfail(
+    reason="""
+During equilibrium rewriter this error happens. Probably one of the routines in SVGD is problematic.
+
+TypeError: The broadcast pattern of the output of scan
+(Matrix(float64, shape=(?, 1))) is inconsistent with the one provided in `output_info`
+(Vector(float64, shape=(?,))). The output on axis 0 is `True`, but it is `False` on axis
+1 in `output_info`. This can happen if one of the dimension is fixed to 1 in the input,
+while it is still variable in the output, or vice-verca. You have to make them consistent,
+e.g. using pytensor.tensor.{unbroadcast, specify_broadcastable}.
+
+Instead of fixing this error it makes sense to rework the internals of the variational to utilize
+pytensor vectorize instead of scan.
+"""
+)
 def test_vi_sampling_jax_svgd():
     with pm.Model():
         x = pm.Normal("x")