Filter warning for batched_dot until we change it (#455)

Author: zaxtax

pymc_extras/inference/pathfinder/pathfinder.py

@@ -502,7 +502,9 @@ def bfgs_sample_dense(
 
     logdet = 2.0 * pt.sum(pt.log(pt.abs(pt.diagonal(Lchol, axis1=-2, axis2=-1))), axis=-1)
 
-    mu = x - pt.batched_dot(H_inv, g)
+    with _warnings.catch_warnings():
+        _warnings.simplefilter("ignore", category=FutureWarning)
+        mu = x - pt.batched_dot(H_inv, g)
 
     phi = pt.matrix_transpose(
         # (L, N, 1)
@@ -572,14 +574,16 @@ def bfgs_sample_sparse(
     logdet += pt.sum(pt.log(alpha), axis=-1)
 
     # NOTE: changed the sign from "x + " to "x -" of the expression to match Stan which differs from Zhang et al., (2022). same for dense version.
-    mu = x - (
-        # (L, N), (L, N) -> (L, N)
-        pt.batched_dot(alpha_diag, g)
-        # beta @ gamma @ beta.T
-        # (L, N, 2J), (L, 2J, 2J), (L, 2J, N) -> (L, N, N)
-        # (L, N, N), (L, N) -> (L, N)
-        + pt.batched_dot((beta @ gamma @ pt.matrix_transpose(beta)), g)
-    )
+    with _warnings.catch_warnings():
+        _warnings.simplefilter("ignore", category=FutureWarning)
+        mu = x - (
+            # (L, N), (L, N) -> (L, N)
+            pt.batched_dot(alpha_diag, g)
+            # beta @ gamma @ beta.T
+            # (L, N, 2J), (L, 2J, 2J), (L, 2J, N) -> (L, N, N)
+            # (L, N, N), (L, N) -> (L, N)
+            + pt.batched_dot((beta @ gamma @ pt.matrix_transpose(beta)), g)
+        )
 
     phi = pt.matrix_transpose(
         # (L, N, 1)

tests/test_pathfinder.py

@@ -18,7 +18,9 @@
 import pymc as pm
 import pytest
 
-pytestmark = pytest.mark.filterwarnings("ignore:compile_pymc was renamed to compile:FutureWarning")
+pytestmark = pytest.mark.filterwarnings(
+    "ignore:compile_pymc was renamed to compile:FutureWarning",
+)
 
 import pymc_extras as pmx