feat: implement PyMC-based Pathfinder VI backend

Add a PyMC/PyTensor implementation of Pathfinder VI as an alternative to the existing BlackJAX backend. Key changes include:

- Implement core Pathfinder components using PyTensor with batched operations
- Add inference_backend parameter to select between PyMC and BlackJAX implementations
- Enable jittering of initial points for Pathfinder

Author: aphc14

pymc_experimental/inference/lbfgs.py

@@ -2,16 +2,14 @@
 from typing import NamedTuple
 
 import numpy as np
-import pytensor.tensor as pt
 
-from pytensor.tensor.variable import TensorVariable
 from scipy.optimize import fmin_l_bfgs_b
 
 
 class LBFGSHistory(NamedTuple):
-    x: TensorVariable
-    f: TensorVariable
-    g: TensorVariable
+    x: np.ndarray
+    f: np.ndarray
+    g: np.ndarray
 
 
 class LBFGSHistoryManager:
@@ -41,9 +39,9 @@ def get_history(self):
         f = self.f_history[: self.count]
         g = self.g_history[: self.count] if self.g_history is not None else None
         return LBFGSHistory(
-            x=pt.as_tensor(x, dtype="float64"),
-            f=pt.as_tensor(f, dtype="float64"),
-            g=pt.as_tensor(g, dtype="float64"),
+            x=x,
+            f=f,
+            g=g,
         )
 
     def __call__(self, x):

pymc_experimental/inference/pathfinder.py

@@ -34,6 +34,7 @@
 from pymc.model.core import Point
 from pymc.sampling.jax import get_jaxified_graph
 from pymc.util import RandomSeed, _get_seeds_per_chain, get_default_varnames
+from pytensor.graph import Apply, Op
 
 from pymc_experimental.inference.lbfgs import lbfgs
 
@@ -311,7 +312,8 @@ def bfgs_sample(
     alpha,
     beta,
     gamma,
-    random_seed: RandomSeed | None = None,
+    rng,
+    # random_seed: RandomSeed | None = None,
 ):
     # batch: L = 8
     # alpha_l: (N,)         => (L, N)
@@ -324,8 +326,6 @@ def bfgs_sample(
     # logdensity: (M,)      => (L, M)
     # theta: (J, N)
 
-    rng = pytensor.shared(np.random.default_rng(seed=random_seed))
-
     if not _batched(x, g, alpha, beta, gamma):
         x = pt.atleast_2d(x)
         g = pt.atleast_2d(g)
@@ -371,6 +371,24 @@ def bfgs_sample(
     return phi, logdensity
 
 
+class LogLike(Op):
+    def __init__(self, logp_func):
+        self.logp_func = logp_func
+        super().__init__()
+
+    def make_node(self, phi_node):
+        # Convert inputs to tensor variables
+        phi_node = pt.as_tensor(phi_node)
+        output_type = pt.tensor(dtype=phi_node.dtype, shape=(None, None))
+        return Apply(self, [phi_node], [output_type])
+
+    def perform(self, node: Apply, phi_node, outputs) -> None:
+        phi_node = phi_node[0]
+        logp_node = np.apply_along_axis(self.logp_func, axis=-1, arr=phi_node)
+        # outputs[0][0] = np.asarray(logp)
+        outputs[0][0] = logp_node
+
+
 def _pymc_pathfinder(
     model,
     x0: np.float64,
@@ -406,38 +424,43 @@ def neg_dlogp_func(x):
         gtol=gtol,
         maxls=maxls,
     )
+    x = pytensor.shared(history.x, "x")
+    g = pytensor.shared(history.g, "g")
 
-    alpha, update_mask = alpha_recover(history.x, history.g)
-
-    beta, gamma = inverse_hessian_factors(alpha, history.x, history.g, update_mask, J=maxcor)
+    alpha, update_mask = alpha_recover(x, g)
 
-    phi, logq_phi = bfgs_sample(
+    beta, gamma = inverse_hessian_factors(alpha, x, g, update_mask, J=maxcor)
+    rng = pytensor.shared(np.random.default_rng(seed=pathfinder_seed))
+    _phi, _logq_phi = bfgs_sample(
         num_samples=num_elbo_draws,
-        x=history.x,
-        g=history.g,
+        x=x,
+        g=g,
         alpha=alpha,
         beta=beta,
         gamma=gamma,
-        random_seed=pathfinder_seed,
+        rng=rng,
     )
+    sample_phi_fn = pytensor.function([alpha, beta, gamma], [_phi, _logq_phi])
+    phi, logq_phi = sample_phi_fn(alpha.eval(), beta.eval(), gamma.eval())
 
     # .vectorize is slower than apply_along_axis
-    logp_phi = np.apply_along_axis(logp_func, axis=-1, arr=phi.eval())
-    logq_phi = logq_phi.eval()
-    elbo = (logp_phi - logq_phi).mean(axis=-1)
-    lstar = np.argmax(elbo)
+    loglike = LogLike(logp_func)
+    logp_phi = loglike(phi)
+    elbo = pt.mean(logp_phi - logq_phi, axis=-1)
+    l_star = pt.argmax(elbo)
 
+    rng.set_value(np.random.default_rng(seed=sample_seed))
     psi, logq_psi = bfgs_sample(
         num_samples=num_draws,
-        x=history.x[lstar],
-        g=history.g[lstar],
-        alpha=alpha[lstar],
-        beta=beta[lstar],
-        gamma=gamma[lstar],
-        random_seed=sample_seed,
+        x=x[l_star],
+        g=g[l_star],
+        alpha=alpha[l_star],
+        beta=beta[l_star],
+        gamma=gamma[l_star],
+        rng=rng,
     )
 
-    return psi[0].eval(), logq_psi, logp_func
+    return psi[0].eval(), logq_psi.eval()
 
 
 def fit_pathfinder(
@@ -492,7 +515,7 @@ def fit_pathfinder(
 
     # TODO: make better
     if inference_backend == "pymc":
-        pathfinder_samples, logq_psi, logp_func = _pymc_pathfinder(
+        pathfinder_samples, logq_psi = _pymc_pathfinder(
             model,
             ip,
             maxcor=maxcor,