Fix failed test from PR pymc-devs#443 and replaced batched_dot for einsum

* Modified the calculation of  in  and  to use  instead of the deprecated  function.
* Added calculations for  in  to improve readability.
* Renamed  to  in pathfinder functions to reflect API changes.
* Updated warning filters to ignore UserWarnings related to einsum subscripts.

Author: aphc14

pymc_extras/inference/pathfinder/pathfinder.py

@@ -40,7 +40,7 @@
 from pymc.model import modelcontext
 from pymc.model.core import Point
 from pymc.pytensorf import (
-    compile_pymc,
+    compile,
     find_rng_nodes,
     reseed_rngs,
 )
@@ -77,7 +77,9 @@
 
 logger = logging.getLogger(__name__)
 _warnings.filterwarnings(
-    "ignore", category=FutureWarning, message="compile_pymc was renamed to compile"
+    "ignore",
+    category=UserWarning,
+    message="The same einsum subscript is used for a broadcastable and non-broadcastable dimension",
 )
 
 REGULARISATION_TERM = 1e-8
@@ -142,7 +144,7 @@ def get_logp_dlogp_of_ravel_inputs(
         [model.logp(jacobian=jacobian), model.dlogp(jacobian=jacobian)],
         model.value_vars,
     )
-    logp_dlogp_fn = compile_pymc([inputs], (logP, dlogP), **compile_kwargs)
+    logp_dlogp_fn = compile([inputs], (logP, dlogP), **compile_kwargs)
     logp_dlogp_fn.trust_input = True
 
     return logp_dlogp_fn
@@ -502,7 +504,7 @@ 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)
+    mu = x - pt.einsum("ijk,ik->ij", H_inv, g)
 
     phi = pt.matrix_transpose(
         # (L, N, 1)
@@ -571,15 +573,12 @@ def bfgs_sample_sparse(
     logdet = 2.0 * pt.sum(pt.log(pt.abs(pt.diagonal(Lchol, axis1=-2, axis2=-1))), axis=-1)
     logdet += pt.sum(pt.log(alpha), axis=-1)
 
+    # inverse Hessian
+    # (L, N, N) + (L, N, 2J), (L, 2J, 2J), (L, 2J, N) -> (L, N, N)
+    H_inv = alpha_diag + (beta @ gamma @ pt.matrix_transpose(beta))
+
     # 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)
-    )
+    mu = x - pt.einsum("ijk,ik->ij", H_inv, g)
 
     phi = pt.matrix_transpose(
         # (L, N, 1)
@@ -853,7 +852,7 @@ def make_pathfinder_body(
 
     # return psi, logP_psi, logQ_psi, elbo_argmax
 
-    pathfinder_body_fn = compile_pymc(
+    pathfinder_body_fn = compile(
         [x_full, g_full],
         [psi, logP_psi, logQ_psi, elbo_argmax],
         **compile_kwargs,

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:The same einsum subscript is used for a broadcastable and non-broadcastable dimension:UserWarning"
+)
 
 import pymc_extras as pmx