#   Copyright 2022 The PyMC Developers
#
#   Licensed under the Apache License, Version 2.0 (the "License");
#   you may not use this file except in compliance with the License.
#   You may obtain a copy of the License at
#
#       http://www.apache.org/licenses/LICENSE-2.0
#
#   Unless required by applicable law or agreed to in writing, software
#   distributed under the License is distributed on an "AS IS" BASIS,
#   WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
#   See the License for the specific language governing permissions and
#   limitations under the License.


import collections
import contextlib
import logging
import multiprocessing as mp
import time
import warnings

from collections import Counter
from collections.abc import Callable, Iterator
from dataclasses import asdict, dataclass, field, replace
from enum import Enum, auto
from typing import Any, Literal, Self, TypeAlias

import arviz as az
import numpy as np
import pymc as pm
import pytensor
import pytensor.tensor as pt

from numpy.typing import NDArray
from packaging import version
from pymc import Model
from pymc.backends.arviz import coords_and_dims_for_inferencedata
from pymc.blocking import DictToArrayBijection, RaveledVars
from pymc.initial_point import make_initial_point_fn
from pymc.model import modelcontext
from pymc.model.core import Point
from pymc.progress_bar import CustomProgress, default_progress_theme
from pymc.pytensorf import compile
from pymc.sampling.mcmc import setup_cores_blas_cores
from pymc.sampling.parallel import _cpu_count, _initialize_multiprocessing_context
from pymc.util import (
    RandomSeed,
    _get_seeds_per_chain,
    get_default_varnames,
)
from pytensor.compile.function.types import Function
from pytensor.compile.mode import FAST_COMPILE, Mode
from pytensor.graph import clone_replace, vectorize_graph
from pytensor.tensor import TensorVariable
from pytensor.tensor.optimize import LRUCache1
from rich.console import Console, Group
from rich.padding import Padding
from rich.progress import TextColumn, TimeElapsedColumn
from rich.table import Column, Table
from rich.text import Text
from threadpoolctl import threadpool_limits

from pymc_extras.inference.laplace_approx.idata import add_data_to_inference_data
from pymc_extras.inference.pathfinder.importance_sampling import (
    importance_sampling as _importance_sampling,
)
from pymc_extras.inference.pathfinder.lbfgs import (
    LBFGS,
    LBFGSException,
    LBFGSInitFailed,
    LBFGSStatus,
    _check_lbfgs_curvature_condition,
)

logger = logging.getLogger(__name__)

REGULARISATION_TERM = 1e-8
DEFAULT_LINKER = "cvm_nogc"

SinglePathfinderFn: TypeAlias = Callable[[int], "PathfinderResult"]


def get_jaxified_logp_of_ravel_inputs(model: Model, jacobian: bool = True) -> Callable:
    """
    Get a JAX function that computes the log-probability of a PyMC model with ravelled inputs.

    Parameters
    ----------
    model : Model
        PyMC model to compute log-probability and gradient.
    jacobian : bool, optional
        Whether to include the Jacobian in the log-probability computation, by default True. Setting to False (not recommended) may result in very high values for pareto k.

    Returns
    -------
    Function
        A JAX function that computes the log-probability of a PyMC model with ravelled inputs.
    """

    from pymc.sampling.jax import get_jaxified_graph

    # TODO: JAX: test if we should get jaxified graph of dlogp as well
    new_logprob, new_input = pm.pytensorf.join_nonshared_inputs(
        model.initial_point(), (model.logp(jacobian=jacobian),), model.value_vars, ()
    )

    logp_func_list = get_jaxified_graph([new_input], new_logprob)

    def logp_func(x):
        return logp_func_list(x)[0]

    return logp_func


def get_logp_dlogp_of_ravel_inputs(
    model: Model, jacobian: bool = True, **compile_kwargs
) -> Function:
    """
    Get the log-probability and its gradient for a PyMC model with ravelled inputs.

    Parameters
    ----------
    model : Model
        PyMC model to compute log-probability and gradient.
    jacobian : bool, optional
        Whether to include the Jacobian in the log-probability computation, by default True. Setting to False (not recommended) may result in very high values for pareto k.
    **compile_kwargs : dict
        Additional keyword arguments to pass to the compile function.

    Returns
    -------
    Function
        A compiled PyTensor function that computes the log-probability and its gradient given ravelled inputs.
    """

    (logP, dlogP), inputs = pm.pytensorf.join_nonshared_inputs(
        model.initial_point(),
        [model.logp(jacobian=jacobian), model.dlogp(jacobian=jacobian)],
        model.value_vars,
    )
    logp_dlogp_fn = compile([inputs], (logP, dlogP), **compile_kwargs)
    logp_dlogp_fn.trust_input = True

    return logp_dlogp_fn


def get_neg_logp_dlogp_of_ravel_inputs(
    model: Model, jacobian: bool = True, **compile_kwargs
) -> Function:
    """Compiled function (x,) -> (-logP, -dlogP) for L-BFGS minimization."""
    (logP, dlogP), inputs = pm.pytensorf.join_nonshared_inputs(
        model.initial_point(),
        [model.logp(jacobian=jacobian), model.dlogp(jacobian=jacobian)],
        model.value_vars,
    )
    neg_logp_dlogp_fn = compile([inputs], (-logP, -dlogP), **compile_kwargs)
    neg_logp_dlogp_fn.trust_input = True
    return neg_logp_dlogp_fn


def get_batched_logp_of_ravel_inputs(
    model: Model, jacobian: bool = True, vectorize: bool = False, **compile_kwargs
) -> Function:
    """Get a batched logP function: (B, N) -> (B,) evaluated in one compiled call.

    Parameters
    ----------
    model : Model
        PyMC model.
    jacobian : bool, optional
        Whether to include the Jacobian, by default True.
    vectorize : bool, optional
        If True, use vectorize_graph instead of pytensor.map. Default False.
        vectorize_graph may be faster for small models but may use more memory
        for large parameter counts.
    **compile_kwargs : dict
        Additional keyword arguments to pass to compile.

    Returns
    -------
    Function
        Compiled function taking a (B, N) array and returning (B,) log-probabilities.
    """
    (logP,), single_input = pm.pytensorf.join_nonshared_inputs(
        model.initial_point(),
        [model.logp(jacobian=jacobian)],
        model.value_vars,
    )
    batch_input = pt.matrix("batch_input", dtype="float64")  # (B, N)
    if vectorize:
        batched_logP = vectorize_graph(logP, replace={single_input: batch_input})
    else:
        # pytensor.map loops over rows of batch_input one at a time; this avoids
        # broadcasting large intermediate tensors (faster than vectorize_graph for
        # models with many parameters).
        batched_logP, _ = pytensor.map(
            fn=lambda x_i: clone_replace([logP], replace={single_input: x_i})[0],
            sequences=[batch_input],
        )
    batched_fn = compile([batch_input], batched_logP, **compile_kwargs)
    batched_fn.trust_input = True
    return batched_fn


def convert_flat_trace_to_idata(
    samples: NDArray,
    include_transformed: bool = False,
    postprocessing_backend: Literal["cpu", "gpu"] = "cpu",
    inference_backend: Literal["pymc", "blackjax"] = "pymc",
    model: Model | None = None,
    importance_sampling: Literal["psis", "psir", "identity"] | None = "psis",
) -> az.InferenceData:
    """convert flattened samples to arviz InferenceData format.

    Parameters
    ----------
    samples : NDArray
        flattened samples
    include_transformed : bool
        whether to include transformed variables
    postprocessing_backend : str
        backend for postprocessing transformations, either "cpu" or "gpu"
    inference_backend : str
        backend for inference, either "pymc" or "blackjax"
    model : Model | None
        pymc model for variable transformations
    importance_sampling : str
        importance sampling method used, affects input samples shape

    Returns
    -------
    InferenceData
        arviz inference data object
    """

    if importance_sampling is None:
        # samples.ndim == 3 in this case, otherwise ndim == 2
        num_paths, num_pdraws, N = samples.shape
        samples = samples.reshape(-1, N)

    model = modelcontext(model)
    ip = model.initial_point()
    ip_point_map_info = DictToArrayBijection.map(ip).point_map_info
    trace = collections.defaultdict(list)
    for sample in samples:
        raveld_vars = RaveledVars(sample, ip_point_map_info)
        point = DictToArrayBijection.rmap(raveld_vars, ip)
        for p, v in point.items():
            trace[p].append(v.tolist())

    trace = {k: np.asarray(v)[None, ...] for k, v in trace.items()}

    var_names = model.unobserved_value_vars
    vars_to_sample = list(get_default_varnames(var_names, include_transformed=include_transformed))
    logger.info("Transforming variables...")

    if inference_backend == "pymc":
        # vectorize_graph batches over trace dims; output size matches input, no extra intermediates.
        new_shapes = [v.ndim * (None,) for v in trace.values()]
        replace = {
            var: pt.tensor(dtype="float64", shape=new_shapes[i])
            for i, var in enumerate(model.value_vars)
        }

        outputs = vectorize_graph(vars_to_sample, replace=replace)

        fn = pytensor.function(
            inputs=[*list(replace.values())],
            outputs=outputs,
            mode=FAST_COMPILE,
            on_unused_input="ignore",
        )
        fn.trust_input = True
        result = fn(*list(trace.values()))
        if not isinstance(result, list):
            result = [result]

        if importance_sampling is None:
            result = [res.reshape(num_paths, num_pdraws, *res.shape[2:]) for res in result]

    elif inference_backend == "blackjax":
        import jax

        from pymc.sampling.jax import get_jaxified_graph

        jax_fn = get_jaxified_graph(inputs=model.value_vars, outputs=vars_to_sample)
        result = jax.vmap(jax.vmap(jax_fn))(
            *jax.device_put(list(trace.values()), jax.devices(postprocessing_backend)[0])
        )

    trace = {v.name: r for v, r in zip(vars_to_sample, result)}
    coords, dims = coords_and_dims_for_inferencedata(model)
    idata = az.from_dict(trace, dims=dims, coords=coords)

    return idata


def alpha_recover(
    x: TensorVariable, g: TensorVariable
) -> tuple[TensorVariable, TensorVariable, TensorVariable]:
    """compute the diagonal elements of the inverse Hessian at each iterations of L-BFGS and filter updates.

    Parameters
    ----------
    x : TensorVariable
        position array, shape (L+1, N)
    g : TensorVariable
        gradient array, shape (L+1, N)

    Returns
    -------
    alpha : TensorVariable
        diagonal elements of the inverse Hessian at each iteration of L-BFGS, shape (L, N)
    s : TensorVariable
        position differences, shape (L, N)
    z : TensorVariable
        gradient differences, shape (L, N)

    Notes
    -----
    shapes: L=batch_size, N=num_params
    """

    def compute_alpha_l(s_l, z_l, alpha_lm1) -> TensorVariable:
        # alpha_lm1: (N,)
        # s_l: (N,)
        # z_l: (N,)
        a = pt.sum(alpha_lm1 * z_l**2)
        b = pt.sum(z_l * s_l)
        c = pt.sum(s_l**2 / alpha_lm1)
        inv_alpha_l = (
            a / (b * alpha_lm1)
            + z_l ** 2 / b
            - (a * s_l ** 2) / (b * c * alpha_lm1**2)
        )  # fmt:off
        return 1.0 / inv_alpha_l

    Lp1, N = x.shape
    s = pt.diff(x, axis=0)
    z = pt.diff(g, axis=0)
    alpha_l_init = pt.ones(N)

    alpha = pytensor.scan(
        fn=compute_alpha_l,
        outputs_info=alpha_l_init,
        sequences=[s, z],
        n_steps=Lp1 - 1,
        allow_gc=False,
        return_updates=False,
    )

    # alpha: (L, N)
    return alpha, s, z


def alpha_step_numpy(alpha_prev: NDArray, s: NDArray, z: NDArray) -> NDArray:
    """Pure-numpy single-step alpha update. Stays in sync with compute_alpha_l in alpha_recover.

    Parameters
    ----------
    alpha_prev : NDArray
        previous alpha vector, shape (N,)
    s : NDArray
        position diff x[l] - x[l-1], shape (N,)
    z : NDArray
        gradient diff g[l] - g[l-1], shape (N,)

    Returns
    -------
    NDArray
        updated alpha, shape (N,)
    """
    a = np.sum(alpha_prev * z**2)
    b = np.sum(z * s)
    c = np.sum(s**2 / alpha_prev)
    # Guard: when b≈0 or c≈0 the formula divides by zero; keep alpha_prev to avoid NaN propagation
    z_sq = float(np.sum(z**2)) + 1e-30
    if abs(b) < 1e-14 * z_sq or c <= 0 or not np.isfinite(c):
        return alpha_prev.copy()
    inv_alpha = (
        a / (b * alpha_prev)
        + z**2 / b
        - (a * s**2) / (b * c * alpha_prev**2)
    )  # fmt: off
    alpha_out = 1.0 / inv_alpha
    if not np.all(np.isfinite(alpha_out)) or np.any(alpha_out <= 0):
        return alpha_prev.copy()
    return alpha_out


def _bfgs_sample_pt(
    x: TensorVariable,
    g: TensorVariable,
    alpha: TensorVariable,
    S: TensorVariable,
    Z: TensorVariable,
    u: TensorVariable,
    J: int,
    N: int,
) -> tuple[TensorVariable, TensorVariable, TensorVariable]:
    """Symbolic L-BFGS inverse-Hessian sample.

    The dense vs sparse path is selected at graph-construction time from
    the compile-time constants N and J, so only one branch is ever compiled.

    Parameters
    ----------
    x : (N,) position
    g : (N,) gradient
    alpha : (N,) diagonal scaling
    S : (N, J) position-diff ring buffer
    Z : (N, J) gradient-diff ring buffer
    u : (M, N) standard-normal draws — M is a dynamic runtime dimension
    J : int, L-BFGS history size (compile-time constant)
    N : int, number of parameters (compile-time constant)

    Returns
    -------
    phi : (M, N) samples
    logQ : (M,) log-density under the approximation
    inv_hessian_diag : (N,) diagonal of sampler covariance used for `phi`
    """
    J2 = 2 * J

    # Inverse Hessian factors
    E = pt.triu(S.T @ Z) + pt.eye(J) * REGULARISATION_TERM  # (J, J)
    eta = pt.diag(E)  # (J,)
    AZ = alpha[:, None] * Z  # (N, J)
    beta = pt.concatenate([AZ, S], axis=1)  # (N, 2J)
    E_inv = pt.linalg.solve_triangular(E, pt.eye(J), lower=False)  # (J, J)
    block_dd = E_inv.T @ (pt.diag(eta) + Z.T @ AZ) @ E_inv  # (J, J)
    top = pt.concatenate([pt.zeros((J, J)), -E_inv], axis=1)  # (J, 2J)
    bot = pt.concatenate([-E_inv.T, block_dd], axis=1)  # (J, 2J)
    gamma = pt.concatenate([top, bot], axis=0)  # (2J, 2J)

    inv_sqrt_alpha = 1.0 / pt.sqrt(alpha)
    sqrt_alpha = pt.sqrt(alpha)

    if J2 >= N:
        # Dense path: form H_inv explicitly (small-N regime, O(N²) is fine)
        isa_d = pt.diag(inv_sqrt_alpha)  # (N, N)
        sa_d = pt.diag(sqrt_alpha)  # (N, N)
        IdN = pt.eye(N) * (1.0 + REGULARISATION_TERM)
        H_inv = sa_d @ (IdN + isa_d @ beta @ gamma @ beta.T @ isa_d) @ sa_d
        Lchol = pt.linalg.cholesky(H_inv, lower=False)  # (N, N)
        logdet = 2.0 * pt.sum(pt.log(pt.abs(pt.diag(Lchol))))
        mu = x - H_inv @ g
        phi = (mu[:, None] + Lchol @ u.T).T  # (M, N)
        inv_hessian_diag = pt.diag(H_inv)
    else:
        # Sparse path: economy QR avoids O(N²) matrices (large-N regime)
        Q, R = pt.linalg.qr(beta * inv_sqrt_alpha[:, None], mode="reduced")  # Q:(N,2J), R:(2J,2J)
        I2J = pt.eye(J2) * (1.0 + REGULARISATION_TERM)
        Lchol = pt.linalg.cholesky(I2J + R @ gamma @ R.T, lower=False)  # (2J, 2J)
        logdet = 2.0 * pt.sum(pt.log(pt.abs(pt.diag(Lchol)))) + pt.sum(pt.log(alpha))
        btg = beta.T @ g[:, None]  # (2J, 1)
        mu = x - ((alpha * g)[:, None] + beta @ (gamma @ btg))[:, 0]  # (N,)
        QtU = Q.T @ u.T  # (2J, M)
        phi = (mu[:, None] + sqrt_alpha[:, None] * (Q @ ((Lchol - pt.eye(J2)) @ QtU) + u.T)).T
        # diag(Q (L L^T - I) Q^T) = row_norm_sq(Q L) - row_norm_sq(Q)
        ql = Q @ Lchol
        inv_hessian_diag = alpha * (1.0 + pt.sum(ql * ql, axis=1) - pt.sum(Q * Q, axis=1))

    logQ = -0.5 * (logdet + pt.sum(u * u, axis=-1) + N * np.log(2.0 * np.pi))
    return phi, logQ, inv_hessian_diag


def make_pathfinder_sample_fn(
    model: Model,
    N: int,
    J: int,
    jacobian: bool,
    compile_kwargs: dict,
    vectorize: bool = False,
    return_inv_hessian_diag: bool = False,
) -> Function:
    """Compile a single PyTensor function covering bfgs sample + batched logP evaluation.

    S and Z are passed as explicit inputs (not shared variables) so each caller
    can pass its own arrays, avoiding shared mutable state.

    The number of draws M is a dynamic runtime dimension — the same compiled
    function handles both ELBO estimation (small M) and final sampling (large M).

    Parameters
    ----------
    model : Model
    N : int, number of unconstrained parameters
    J : int, L-BFGS history size (maxcor)
    jacobian : bool
    compile_kwargs : dict
    vectorize : bool, optional
        If True, use vectorize_graph instead of pytensor.map. Default False.
    return_inv_hessian_diag : bool, optional
        If True, include inverse-Hessian diagonal as 4th output.

    Returns
    -------
    fn : Function
        Compiled: (x, g, alpha, s_win, z_win, u) → (phi, logQ, logP[, inv_hessian_diag])
        where s_win, z_win are (N, J), u is (M, N), and M is a dynamic dimension.
    """
    (logP_single,), single_input = pm.pytensorf.join_nonshared_inputs(
        model.initial_point(),
        [model.logp(jacobian=jacobian)],
        model.value_vars,
    )

    x_sym = pt.vector("x", dtype="float64")
    g_sym = pt.vector("g", dtype="float64")
    alpha_sym = pt.vector("alpha", dtype="float64")
    s_win_sym = pt.matrix("s_win", dtype="float64")  # (N, J)
    z_win_sym = pt.matrix("z_win", dtype="float64")  # (N, J)
    u_sym = pt.matrix("u", dtype="float64")  # (M, N) — M is dynamic

    phi_sym, logQ_sym, inv_hessian_diag_sym = _bfgs_sample_pt(
        x_sym, g_sym, alpha_sym, s_win_sym, z_win_sym, u_sym, J, N
    )

    if vectorize:
        batched_logP_sym = vectorize_graph(logP_single, replace={single_input: phi_sym})
    else:
        batched_logP_sym = pytensor.map(
            fn=lambda x_i: clone_replace([logP_single], replace={single_input: x_i})[0],
            sequences=[phi_sym],
            return_updates=False,
        )

    outputs = [phi_sym, logQ_sym, batched_logP_sym]
    if return_inv_hessian_diag:
        outputs.append(inv_hessian_diag_sym)

    fn = pytensor.function(
        [
            pytensor.In(x_sym, borrow=True),
            pytensor.In(g_sym, borrow=True),
            pytensor.In(alpha_sym, borrow=True),
            pytensor.In(s_win_sym, borrow=True),
            pytensor.In(z_win_sym, borrow=True),
            pytensor.In(u_sym, borrow=True),
        ],
        outputs,
        **compile_kwargs,
    )
    fn.trust_input = True
    return fn


class PathStatus(Enum):
    """
    Statuses of a single-path pathfinder.
    """

    SUCCESS = auto()
    ELBO_ARGMAX_AT_ZERO = auto()
    # Statuses that lead to Exceptions:
    INVALID_LOGP = auto()
    INVALID_LOGQ = auto()
    LBFGS_FAILED = auto()
    PATH_FAILED = auto()
    SINGLE_STEP = auto()


FAILED_PATH_STATUS = [
    PathStatus.INVALID_LOGP,
    PathStatus.INVALID_LOGQ,
    PathStatus.LBFGS_FAILED,
    PathStatus.PATH_FAILED,
    PathStatus.SINGLE_STEP,
]


class PathException(Exception):
    """
    raises a PathException if the path failed.
    """

    DEFAULT_MESSAGE = "Path failed."

    def __init__(self, message=None, status: PathStatus = PathStatus.PATH_FAILED) -> None:
        super().__init__(message or self.DEFAULT_MESSAGE)
        self.status = status


class PathInvalidLogP(PathException):
    """
    raises a PathException if all the logP values in a path are not finite.
    """

    DEFAULT_MESSAGE = "Path failed because all the logP values in a path are not finite."

    def __init__(self, message=None) -> None:
        super().__init__(message or self.DEFAULT_MESSAGE, PathStatus.INVALID_LOGP)


class PathInvalidLogQ(PathException):
    """
    raises a PathException if all the logQ values in a path are not finite.
    """

    DEFAULT_MESSAGE = "Path failed because all the logQ values in a path are not finite."

    def __init__(self, message=None) -> None:
        super().__init__(message or self.DEFAULT_MESSAGE, PathStatus.INVALID_LOGQ)


class SingleStepPathException(PathException):
    """
    raises when the path has only one LBFGS step (insufficient for valid approximation).
    """

    DEFAULT_MESSAGE = "Path failed because only a single step was performed."

    def __init__(self, message=None) -> None:
        super().__init__(message or self.DEFAULT_MESSAGE, PathStatus.SINGLE_STEP)


class LBFGSStreamingCallback:
    """Streaming LBFGS callback: computes ELBO at each accepted step, O(J*N + M*N) peak memory.

    Instead of collecting the full (L+1, N) history, it processes each accepted step
    immediately and tracks only the best state seen so far.

    Parameters
    ----------
    value_grad_fn : Callable
        Single-entry cached value/gradient function (e.g. LRUCache1 from pytensor.tensor.optimize).
    x0 : NDArray
        Initial position, shape (N,).
    sample_logp_fn : Callable
        Compiled PyTensor function (x, g, alpha, s_win, z_win, u) → (phi, logQ, logP).
        Built by make_pathfinder_sample_fn.
    num_elbo_draws : int
        Number of draws per step for ELBO estimation.
    rng : np.random.Generator
        Random number generator for draw generation.
    J : int
        L-BFGS history size (maxcor).
    epsilon : float
        Tolerance for the LBFGS update condition.
    progress_callback : Callable | None
        Optional progress reporting.
    on_step_callback : Callable | None
        If set, called after each accepted step with (x, g, alpha, s_win, z_win, elbo).
        Used by fixture generation to record per-step state.
    """

    def __init__(
        self,
        value_grad_fn: Callable,
        x0: NDArray,
        sample_logp_fn: Callable,
        num_elbo_draws: int,
        rng: np.random.Generator,
        J: int,
        epsilon: float,
        progress_callback: Callable | None = None,
        on_step_callback: Callable | None = None,
    ) -> None:
        self.value_grad_fn = value_grad_fn
        self.sample_logp_fn = sample_logp_fn
        self.num_elbo_draws = num_elbo_draws
        self._rng = rng
        self.J = J
        self.epsilon = epsilon
        self.progress_callback = progress_callback
        self.on_step_callback = on_step_callback

        N = x0.shape[0]
        self._N = N
        _, g0 = value_grad_fn(x0)

        self.x_prev: NDArray = x0.copy()
        self.g_prev: NDArray = np.array(g0, dtype=np.float64)
        self.alpha_prev: NDArray = np.ones(N, dtype=np.float64)

        # Ring buffer: numpy arrays passed as inputs to sample_logp_fn each call.
        # Thread-safe: no shared mutable state across concurrent invocations.
        self.s_win: NDArray = np.zeros((N, J), dtype=np.float64)
        self.z_win: NDArray = np.zeros((N, J), dtype=np.float64)
        self.win_idx: int = -1
        self.best_elbo: float = -np.inf
        self.best_state: dict = {}
        self.best_step_idx: int = 0
        self.step_count: int = 0
        self.any_valid: bool = False
        self.current_elbo: float | None = None
        self._start_time: float = time.time()

    def __call__(self, x: NDArray) -> float | None:
        """Process one accepted LBFGS step. Returns current_elbo for testability."""
        value, g = self.value_grad_fn(x)

        s = x - self.x_prev
        z = g - self.g_prev

        if not (np.all(np.isfinite(g)) and np.isfinite(value)):
            self.current_elbo = None
            return None
        if not _check_lbfgs_curvature_condition(s, z, self.epsilon):
            self.current_elbo = None
            return None

        alpha = alpha_step_numpy(self.alpha_prev, s, z)

        # Ring-buffer update (numpy, O(N))
        self.win_idx = (self.win_idx + 1) % self.J
        self.s_win[:, self.win_idx] = s
        self.z_win[:, self.win_idx] = z

        # Sample + logP in a single compiled call. Pass s_win/z_win as inputs.
        u = self._rng.standard_normal((self.num_elbo_draws, self._N))
        try:
            sample_out = self.sample_logp_fn(x, g, alpha, self.s_win, self.z_win, u)
            _, logQ, logP = sample_out[:3]
            logP = np.asarray(logP)
            logQ = np.asarray(logQ)
            finite = np.isfinite(logP)
            if not np.any(finite):
                elbo = -np.inf
            else:
                logP_safe = np.where(finite, logP, -np.inf)
                elbo = float(np.mean(logP_safe - logQ))
                if not np.isfinite(elbo):
                    elbo = -np.inf
        except Exception:
            elbo = -np.inf

        if np.isfinite(elbo):
            self.any_valid = True

        if self.on_step_callback is not None:
            self.on_step_callback(x, g, alpha, self.s_win.copy(), self.z_win.copy(), elbo)

        if elbo > self.best_elbo:
            self.best_elbo = elbo
            self.best_state = {
                "alpha": alpha.copy(),
                "s_win": self.s_win.copy(),
                "z_win": self.z_win.copy(),
                "win_idx": self.win_idx,
                "x": x.copy(),
                "g": g.copy(),
            }
            self.best_step_idx = self.step_count

        self.alpha_prev = alpha
        self.x_prev = x.copy()
        self.g_prev = g.copy()
        self.step_count += 1

        self.current_elbo = elbo if np.isfinite(elbo) else None

        if self.progress_callback is not None:
            best_elbo = self.best_elbo if np.isfinite(self.best_elbo) else None
            current_elbo = self.current_elbo
            elapsed = time.time() - self._start_time
            steps_per_sec = self.step_count / elapsed if elapsed > 0 else None
            step_size = float(np.linalg.norm(s))
            self.progress_callback(
                {
                    "lbfgs_steps": self.step_count,
                    "best_elbo": best_elbo,
                    "best_ind": self.best_step_idx,
                    "current_elbo": current_elbo,
                    "step_size": step_size,
                    "steps_per_sec": steps_per_sec,
                }
            )

        return elbo


def make_single_pathfinder_fn(
    model,
    num_draws: int,
    maxcor: int | None,
    maxiter: int,
    ftol: float,
    gtol: float,
    maxls: int,
    num_elbo_draws: int,
    jitter: float,
    epsilon: float,
    max_init_retries: int = 10,
    pathfinder_kwargs: dict[str, Any] = {},
    compile_kwargs: dict[str, Any] = {},
) -> SinglePathfinderFn:
    """
    returns a seedable single-path pathfinder function, where it executes a compiled function that performs the local approximation and sampling part of the Pathfinder algorithm.

    Parameters
    ----------
    model : pymc.Model
        The PyMC model to fit the Pathfinder algorithm to.
    num_draws : int
        Number of samples to draw from the single-path approximation.
    maxcor : int | None
        Maximum number of iterations for the L-BFGS optimisation.
    maxiter : int
        Maximum number of iterations for the L-BFGS optimisation.
    ftol : float
        Tolerance for the decrease in the objective function.
    gtol : float
        Tolerance for the norm of the gradient.
    maxls : int
        Maximum number of line search steps for the L-BFGS algorithm.
    num_elbo_draws : int
        Number of draws for the Evidence Lower Bound (ELBO) estimation.
    jitter : float
        Amount of jitter to apply to initial points. Note that Pathfinder may be highly sensitive to the jitter value. It is recommended to increase num_paths when increasing the jitter value.
    epsilon : float
        value used to filter out large changes in the direction of the update gradient at each iteration l in L. Iteration l is only accepted if s·z >= epsilon * ||z||² for each l in L. Matches Zhang et al. (2022) Algorithm 3 with default epsilon=1e-12.
    max_init_retries : int, optional
        Maximum number of re-jitter retries when LBFGSInitFailed is raised (i.e. the initial
        point yields non-finite value/gradient or the first step is rejected). Each retry uses
        a different jitter seed. Defaults to 10.
    pathfinder_kwargs : dict
        Additional keyword arguments for the Pathfinder algorithm. Supported keys:
        ``jacobian`` (bool, default True), ``vectorize`` (bool, default False).
        If ``vectorize`` is True, use vectorize_graph instead of pytensor.map for
        batched logP evaluation; may be faster for small models.
    compile_kwargs : dict
        Additional keyword arguments for the PyTensor compiler. If not provided, a
        performant default is used.

    Returns
    -------
    single_pathfinder_fn : Callable
        A seedable single-path pathfinder function that accepts ``(random_seed, progress_callback=None)``.
    """

    compile_kwargs = {"mode": Mode(linker=DEFAULT_LINKER), **compile_kwargs}
    jacobian = pathfinder_kwargs.get("jacobian", True)
    vectorize = pathfinder_kwargs.get("vectorize", False)
    logp_dlogp_kwargs = {"jacobian": jacobian, **compile_kwargs}

    neg_logp_dlogp_func = get_neg_logp_dlogp_of_ravel_inputs(model, **logp_dlogp_kwargs)

    # initial point
    # TODO: remove make_initial_points function when feature request is implemented: https://github.com/pymc-devs/pymc/issues/7555
    ipfn = make_initial_point_fn(model=model)
    ip = Point(ipfn(None), model=model)
    x_base = DictToArrayBijection.map(ip).data
    N = x_base.shape[0]

    sample_logp_func = make_pathfinder_sample_fn(
        model,
        N,
        maxcor,
        jacobian=jacobian,
        compile_kwargs=compile_kwargs,
        vectorize=vectorize,
    )
    final_sample_logp_func = make_pathfinder_sample_fn(
        model,
        N,
        maxcor,
        jacobian=jacobian,
        compile_kwargs=compile_kwargs,
        vectorize=vectorize,
        return_inv_hessian_diag=True,
    )

    def _check_lbfgs_status(status):
        if status in {LBFGSStatus.INIT_FAILED, LBFGSStatus.INIT_FAILED_LOW_UPDATE_PCT}:
            raise LBFGSInitFailed(status)
        elif status == LBFGSStatus.LBFGS_FAILED:
            raise LBFGSException()

    def _make_result(
        psi, logP_psi, logQ_psi, lbfgs_niter, elbo_argmax, inv_hessian_diag, lbfgs_status
    ):
        if np.all(~np.isfinite(logQ_psi)):
            raise PathInvalidLogQ()
        path_status = PathStatus.ELBO_ARGMAX_AT_ZERO if elbo_argmax == 0 else PathStatus.SUCCESS
        return PathfinderResult(
            samples=psi,
            logP=logP_psi,
            logQ=logQ_psi,
            lbfgs_niter=lbfgs_niter,
            elbo_argmax=elbo_argmax,
            inv_hessian_diag=inv_hessian_diag,
            lbfgs_status=lbfgs_status,
            path_status=path_status,
        )

    def single_pathfinder_fn(
        random_seed: int, progress_callback: Callable | None = None
    ) -> PathfinderResult:
        if progress_callback is not None:
            progress_callback({"status": "running"})

        # Per-path independent copies of compiled functions for process safety.
        local_neg_logp_dlogp = neg_logp_dlogp_func.copy(share_memory=False)
        cached_fn = LRUCache1(local_neg_logp_dlogp, copy_x=True)

        local_lbfgs = LBFGS(cached_fn, maxcor, maxiter, ftol, gtol, maxls, epsilon)

        local_sample_logp = sample_logp_func.copy(share_memory=False)
        local_final_sample_logp = final_sample_logp_func.copy(share_memory=False)

        lbfgs_status = LBFGSStatus.LBFGS_FAILED  # default before LBFGS runs
        try:
            # Derive base seeds once from random_seed. init_seed is an int so we can
            # safely offset it per retry attempt (avoids Generator arithmetic issues).
            _base_init, elbo_seed, final_seed = _get_seeds_per_chain(random_seed, 3)

            for attempt in range(max_init_retries + 1):
                try:
                    init_seed = _base_init + attempt  # different jitter per retry
                    rng = np.random.default_rng(init_seed)
                    jitter_value = rng.uniform(-jitter, jitter, size=x_base.shape)
                    x0 = x_base + jitter_value

                    # Fresh NumPy RNG each attempt → reproducible regardless
                    # of how many ELBO steps the previous attempt took.
                    elbo_rng = np.random.default_rng(elbo_seed + attempt)
                    cached_fn.clear_cache()

                    streaming_cb = LBFGSStreamingCallback(
                        value_grad_fn=cached_fn,
                        x0=x0,
                        sample_logp_fn=local_sample_logp,
                        num_elbo_draws=num_elbo_draws,
                        rng=elbo_rng,
                        J=maxcor,
                        epsilon=epsilon,
                        progress_callback=progress_callback,
                    )

                    lbfgs_niter, lbfgs_status = local_lbfgs.minimize_streaming(streaming_cb, x0)
                    _check_lbfgs_status(lbfgs_status)

                    if lbfgs_niter < 2:
                        raise SingleStepPathException()

                    if not streaming_cb.any_valid:
                        raise PathInvalidLogP()

                    elbo_argmax = streaming_cb.best_step_idx
                    best_state = streaming_cb.best_state

                    final_rng = np.random.default_rng(final_seed)
                    u_final = final_rng.standard_normal((num_draws, N))
                    if progress_callback is not None:
                        progress_callback(
                            {"status": "sampling", "current_elbo": None, "step_size": None}
                        )

                    sample_out = local_final_sample_logp(
                        best_state["x"],
                        best_state["g"],
                        best_state["alpha"],
                        best_state["s_win"],
                        best_state["z_win"],
                        u_final,
                    )
                    phi_final, logQ_psi_flat, logP_psi_flat, inv_hessian_diag = sample_out
                    phi_final = np.asarray(phi_final)
                    logQ_psi_flat = np.asarray(logQ_psi_flat)
                    logP_psi_flat = np.asarray(logP_psi_flat)
                    inv_hessian_diag = np.asarray(inv_hessian_diag)[None]
                    # Add batch dim L=1 to match downstream expectations
                    psi = phi_final[None]  # (1, M, N)
                    logP_psi = logP_psi_flat[None]  # (1, M)
                    logQ_psi = logQ_psi_flat[None]  # (1, M)

                    break  # success — exit retry loop

                except (LBFGSInitFailed, SingleStepPathException) as e:
                    if attempt < max_init_retries:
                        logger.debug(
                            "%s on attempt %d/%d, retrying with different jitter...",
                            type(e).__name__,
                            attempt + 1,
                            max_init_retries,
                        )
                        if progress_callback is not None:
                            progress_callback({"status": f"retry {attempt + 1}"})
                    else:
                        if progress_callback is not None:
                            progress_callback(
                                {
                                    "status": "lbfgs_fail"
                                    if isinstance(e, LBFGSInitFailed)
                                    else "failed"
                                }
                            )
                        if isinstance(e, LBFGSInitFailed):
                            return PathfinderResult(
                                lbfgs_status=e.status,
                                path_status=PathStatus.LBFGS_FAILED,
                            )
                        return PathfinderResult(
                            lbfgs_status=lbfgs_status,
                            path_status=PathStatus.SINGLE_STEP,
                        )

            result = _make_result(
                psi,
                logP_psi,
                logQ_psi,
                lbfgs_niter,
                elbo_argmax,
                inv_hessian_diag,
                lbfgs_status,
            )
            if progress_callback is not None:
                status_str = (
                    "elbo@0" if result.path_status == PathStatus.ELBO_ARGMAX_AT_ZERO else "ok"
                )
                progress_callback(
                    {"status": status_str, "lbfgs_steps": int(lbfgs_niter), "best_ind": elbo_argmax}
                )
            return result

        except LBFGSException as e:
            if progress_callback is not None:
                progress_callback({"status": "lbfgs_fail"})
            return PathfinderResult(
                lbfgs_status=e.status,
                path_status=PathStatus.LBFGS_FAILED,
            )
        except PathException as e:
            if progress_callback is not None:
                progress_callback({"status": "failed"})
            return PathfinderResult(
                lbfgs_status=lbfgs_status,
                path_status=e.status,
            )

    return single_pathfinder_fn


def _default_cores(num_paths: int, cores: int | None) -> int:
    """Default cores for parallel pathfinder, mirroring pm.sample."""
    if cores is not None:
        return min(cores, num_paths)
    return min(4, _cpu_count(), num_paths)


class _PipeCallback:
    """Picklable progress callback that relays updates through a multiprocessing Pipe."""

    def __init__(self, conn: Any, idx: int) -> None:
        self.conn = conn
        self.idx = idx

    def __call__(self, info: dict) -> None:
        try:
            self.conn.send((self.idx, info))
        except Exception:
            pass


def _run_path(
    fn_pickled: bytes,
    seed: int,
    path_id: int,
    progress_conn: Any,
    blas_cores: int | None,
    mp_start_method: str,
) -> "PathfinderResult":
    """Worker: unpickle fn, run with threadpool_limits when blas_cores set (non-fork)."""
    import cloudpickle

    from pytensor.compile.compilelock import lock_ctx

    ctx = (
        threadpool_limits(limits=blas_cores)
        if mp_start_method != "fork" and blas_cores is not None
        else contextlib.nullcontext()
    )
    with ctx:
        with lock_ctx(timeout=-1):
            fn = cloudpickle.loads(fn_pickled)
        cb = _PipeCallback(progress_conn, path_id)
        try:
            return fn(seed, cb)
        finally:
            try:
                progress_conn.send((path_id, None))  # sentinel
            except Exception:
                pass


def _execute_concurrently(
    fn: SinglePathfinderFn,
    seeds: list[int],
    cores: int,
    blas_cores: int | None,
    progress_callbacks: list[Callable | None] | None = None,
    mp_ctx: mp.context.BaseContext | str | None = None,
) -> Iterator["PathfinderResult"]:
    """Execute pathfinder runs concurrently via ProcessPoolExecutor.

    Uses Pipe instead of Manager().Queue() to avoid spawn bootstrapping issues
    when mp_ctx is 'spawn' and the main module is still loading.
    """
    import threading

    from concurrent.futures import ProcessPoolExecutor, as_completed

    import cloudpickle

    mp_ctx = _initialize_multiprocessing_context(mp_ctx, quiet=True)
    fn_pickled = cloudpickle.dumps(fn, protocol=-1)

    # One pipe per worker; avoids Manager() which spawns a process and triggers
    # "An attempt has been made to start a new process before bootstrapping" when
    # the main module is still loading (e.g. script run without if __name__ guard).
    n_workers = len(seeds)
    parent_conns = []
    child_conns = []
    for _ in range(n_workers):
        parent, child = mp_ctx.Pipe(duplex=False)
        parent_conns.append(parent)
        child_conns.append(child)

    sentinel_count: list[int] = [0]
    sentinel_lock = threading.Lock()

    def _listener() -> None:
        from multiprocessing.connection import wait

        while True:
            ready = wait(parent_conns, timeout=0.1)
            for conn in ready:
                try:
                    idx, info = conn.recv()
                except EOFError:
                    continue
                if info is None:
                    with sentinel_lock:
                        sentinel_count[0] += 1
                    if sentinel_count[0] >= n_workers:
                        return
                    continue
                if (
                    progress_callbacks
                    and idx < len(progress_callbacks)
                    and progress_callbacks[idx] is not None
                ):
                    progress_callbacks[idx](info)

    listener = threading.Thread(target=_listener, daemon=True)
    listener.start()

    def _run_executor():
        with ProcessPoolExecutor(max_workers=cores, mp_context=mp_ctx) as executor:
            futures = {
                executor.submit(
                    _run_path,
                    fn_pickled,
                    seed,
                    i,
                    child_conns[i],
                    blas_cores,
                    mp_ctx.get_start_method(),
                ): i
                for i, seed in enumerate(seeds)
            }
            for f in as_completed(futures):
                yield f.result()

    try:
        yield from _run_executor()
    except RuntimeError as e:
        if "bootstrapping" in str(e) and mp_ctx.get_start_method() == "spawn":
            raise RuntimeError(
                "Pathfinder with mp_ctx='spawn' requires the entry point to be "
                "guarded with `if __name__ == '__main__':`. When using spawn, "
                "wrap your fit_pathfinder call (or the function that contains it) "
                "in that block. See: "
                "https://docs.python.org/3/library/multiprocessing.html"
                "#the-spawn-and-forkserver-start-methods"
            ) from e
        raise
    finally:
        for c in child_conns:
            try:
                c.close()
            except Exception:
                pass
        listener.join(timeout=5)


def _execute_serially(
    fn: SinglePathfinderFn,
    seeds: list[int],
    progress_callbacks: list[Callable | None] | None = None,
) -> Iterator["PathfinderResult"]:
    """Execute pathfinder runs serially."""
    callbacks = progress_callbacks or [None] * len(seeds)
    for seed, cb in zip(seeds, callbacks):
        yield fn(seed, cb)


def make_generator(
    concurrent: Literal["process"] | None,
    fn: SinglePathfinderFn,
    seeds: list[int],
    cores: int | None = None,
    blas_cores: int | None | Literal["auto"] = "auto",
    progress_callbacks: list[Callable | None] | None = None,
    mp_ctx: mp.context.BaseContext | str | None = None,
) -> Iterator["PathfinderResult"]:
    """Generator for executing pathfinder runs concurrently or serially."""
    if concurrent is not None:
        num_paths = len(seeds)
        effective_cores = _default_cores(num_paths, cores)
        _, _, num_blas_per_worker = setup_cores_blas_cores(
            blas_cores,
            num_paths,
            effective_cores,
            _initialize_multiprocessing_context(mp_ctx, quiet=True),
        )
        yield from _execute_concurrently(
            fn,
            seeds,
            cores=effective_cores,
            blas_cores=num_blas_per_worker,
            progress_callbacks=progress_callbacks,
            mp_ctx=mp_ctx,
        )
    else:
        yield from _execute_serially(fn, seeds, progress_callbacks)


@dataclass(slots=True, frozen=True)
class PathfinderResult:
    """
    container for storing results from a single pathfinder run.

    Attributes
    ----------
        samples: posterior samples (1, M, N)
        logP: log probability of model (1, M)
        logQ: log probability of approximation (1, M)
        lbfgs_niter: number of lbfgs iterations (1,)
        elbo_argmax: elbo values at convergence (1,)
        lbfgs_status: LBFGS status
        path_status: path status

    where:
        M: number of samples
        N: number of parameters
    """

    samples: NDArray | None = None
    logP: NDArray | None = None
    logQ: NDArray | None = None
    lbfgs_niter: NDArray | None = None
    elbo_argmax: NDArray | None = None
    inv_hessian_diag: NDArray | None = None
    lbfgs_status: LBFGSStatus = LBFGSStatus.LBFGS_FAILED
    path_status: PathStatus = PathStatus.PATH_FAILED


@dataclass(frozen=True)
class PathfinderConfig:
    """configuration parameters for a single pathfinder"""

    num_draws: int  # same as num_draws_per_path
    maxcor: int
    maxiter: int
    ftol: float
    gtol: float
    maxls: int
    jitter: float
    epsilon: float
    num_elbo_draws: int


@dataclass(slots=True, frozen=True)
class MultiPathfinderResult:
    """
    container for aggregating results from multiple paths.

    Attributes
    ----------
        samples: posterior samples (S, M, N)
        logP: log probability of model (S, M)
        logQ: log probability of approximation (S, M)
        lbfgs_niter: number of lbfgs iterations (S,)
        elbo_argmax: elbo values at convergence (S,)
        lbfgs_status: counter for LBFGS status occurrences
        path_status: counter for path status occurrences
        importance_sampling: importance sampling method used
        warnings: list of warnings
        pareto_k
        pathfinder_config: pathfinder configuration
        compile_time
        compute_time
    where:
        S: number of successful paths, where S <= num_paths
        M: number of samples per path
        N: number of parameters
    """

    samples: NDArray | None = None
    logP: NDArray | None = None
    logQ: NDArray | None = None
    lbfgs_niter: NDArray | None = None
    elbo_argmax: NDArray | None = None
    inv_hessian_diag: NDArray | None = None
    lbfgs_status: Counter = field(default_factory=Counter)
    path_status: Counter = field(default_factory=Counter)
    importance_sampling: str | None = "psis"
    warnings: list[str] = field(default_factory=list)
    pareto_k: float | None = None

    # config
    num_paths: int | None = None
    num_draws: int | None = None
    pathfinder_config: PathfinderConfig | None = None

    # timing
    compile_time: float | None = None
    compute_time: float | None = None

    all_paths_failed: bool = False  # raises ValueError if all paths failed

    @classmethod
    def from_path_results(cls, path_results: list[PathfinderResult]) -> "MultiPathfinderResult":
        """aggregate successful pathfinder results and count the occurrences of each status in PathStatus and LBFGSStatus"""

        NUMERIC_ATTRIBUTES = [
            "samples",
            "logP",
            "logQ",
            "lbfgs_niter",
            "elbo_argmax",
            "inv_hessian_diag",
        ]

        success_results = []
        mpr = cls()

        for pr in path_results:
            if pr.path_status not in FAILED_PATH_STATUS:
                success_results.append(tuple(getattr(pr, attr) for attr in NUMERIC_ATTRIBUTES))

            mpr.lbfgs_status[pr.lbfgs_status] += 1
            mpr.path_status[pr.path_status] += 1

        warnings = _get_status_warning(mpr)

        if success_results:
            results_arr = [np.asarray(x) for x in zip(*success_results)]
            return cls(
                *[np.concatenate(x) if x.ndim > 1 else x for x in results_arr],
                lbfgs_status=mpr.lbfgs_status,
                path_status=mpr.path_status,
                warnings=warnings,
            )
        else:
            return cls(
                lbfgs_status=mpr.lbfgs_status,
                path_status=mpr.path_status,
                warnings=warnings,
                all_paths_failed=True,  # raises ValueError later
            )

    def with_timing(self, compile_time: float, compute_time: float) -> Self:
        """add timing information"""
        return replace(self, compile_time=compile_time, compute_time=compute_time)

    def with_pathfinder_config(self, config: PathfinderConfig) -> Self:
        """add pathfinder configuration"""
        return replace(self, pathfinder_config=config)

    def with_warnings(self, warnings: list[str]) -> Self:
        """add warnings"""
        return replace(self, warnings=warnings)

    def with_importance_sampling(
        self,
        num_draws: int,
        method: Literal["psis", "psir", "identity"] | None,
        random_seed: int | None = None,
    ) -> Self:
        """perform importance sampling"""
        if not self.all_paths_failed:
            isres = _importance_sampling(
                samples=self.samples,
                logP=self.logP,
                logQ=self.logQ,
                num_draws=num_draws,
                method=method,
                random_seed=random_seed,
            )
            return replace(
                self,
                samples=isres.samples,
                importance_sampling=method,
                warnings=[*self.warnings, *isres.warnings],
                pareto_k=isres.pareto_k,
            )
        else:
            return self

    def create_summary(self) -> Table:
        """create rich table summary of pathfinder results"""
        table = Table(
            title="Pathfinder Results",
            title_style="none",
            title_justify="left",
            show_header=False,
            box=None,
            padding=(0, 2),
            show_edge=False,
        )
        table.add_column("Description")
        table.add_column("Value")

        # model info
        if self.samples is not None:
            table.add_row("")
            table.add_row("No. model parameters", str(self.samples.shape[-1]))

        # config
        if self.pathfinder_config is not None:
            table.add_row("")
            table.add_row("Configuration:")
            table.add_row("num_draws_per_path", str(self.pathfinder_config.num_draws))
            table.add_row("history size (maxcor)", str(self.pathfinder_config.maxcor))
            table.add_row("max iterations", str(self.pathfinder_config.maxiter))
            table.add_row("ftol", f"{self.pathfinder_config.ftol:.2e}")
            table.add_row("gtol", f"{self.pathfinder_config.gtol:.2e}")
            table.add_row("max line search", str(self.pathfinder_config.maxls))
            table.add_row("jitter", f"{self.pathfinder_config.jitter}")
            table.add_row("epsilon", f"{self.pathfinder_config.epsilon:.2e}")
            table.add_row("ELBO draws", str(self.pathfinder_config.num_elbo_draws))

        # lbfgs
        table.add_row("")
        table.add_row("LBFGS Status:")
        for status, count in self.lbfgs_status.items():
            table.add_row(str(status.name), str(count))

        if self.lbfgs_niter is not None:
            table.add_row(
                "L-BFGS iterations",
                f"mean {np.mean(self.lbfgs_niter):.0f} ± std {np.std(self.lbfgs_niter):.0f}",
            )

        # paths
        table.add_row("")
        table.add_row("Path Status:")
        for status, count in self.path_status.items():
            table.add_row(str(status.name), str(count))

        if self.elbo_argmax is not None:
            table.add_row(
                "ELBO argmax",
                f"mean {np.mean(self.elbo_argmax):.0f} ± std {np.std(self.elbo_argmax):.0f}",
            )

        # importance sampling section
        if not self.all_paths_failed:
            table.add_row("")
            table.add_row("Importance Sampling:")
            table.add_row("Method", self.importance_sampling)
            if self.pareto_k is not None:
                table.add_row("Pareto k", f"{self.pareto_k:.2f}")

        if self.compile_time is not None:
            table.add_row("")
            table.add_row("Timing (seconds):")
            table.add_row("Compile", f"{self.compile_time:.2f}")

        if self.compute_time is not None:
            table.add_row("Compute", f"{self.compute_time:.2f}")

        if self.compile_time is not None and self.compute_time is not None:
            table.add_row("Total", f"{self.compile_time + self.compute_time:.2f}")

        return table

    def display_summary(self) -> None:
        """display summary including warnings"""
        console = Console()
        summary = self.create_summary()

        # warning messages
        if self.warnings:
            warning_text = [
                Text(),  # blank line
                Text("Warnings:"),
                *(
                    Padding(
                        Text("- " + warning, no_wrap=False).wrap(console, width=console.width - 6),
                        (0, 0, 0, 2),  # left padding only
                    )
                    for warning in self.warnings
                ),
            ]
            output = Group(summary, *warning_text)
        else:
            output = summary

        console.print(output)


def _get_status_warning(mpr: MultiPathfinderResult) -> list[str]:
    """get list of relevant LBFGSStatus and PathStatus warnings given a MultiPathfinderResult"""
    warnings = []

    lbfgs_status_message = {
        LBFGSStatus.MAX_ITER_REACHED: "MAX_ITER_REACHED: LBFGS maximum number of iterations reached. Consider increasing maxiter if this occurence is high relative to the number of paths.",
        LBFGSStatus.INIT_FAILED: "INIT_FAILED: LBFGS failed to initialize. Consider reparameterizing the model or reducing jitter if this occurence is high relative to the number of paths.",
        LBFGSStatus.NON_FINITE: "NON_FINITE: LBFGS objective function produced inf or nan at the last iteration. Consider reparameterizing the model or adjusting the pathfinder arguments if this occurence is high relative to the number of paths.",
        LBFGSStatus.LOW_UPDATE_PCT: "LOW_UPDATE_PCT: Majority of LBFGS iterations were not accepted due to the either: (1) LBFGS function or gradient values containing too many inf or nan values or (2) gradient changes being significantly large, set by epsilon. Consider reparameterizing the model, adjusting initvals or jitter or other pathfinder arguments if this occurence is high relative to the number of paths.",
        LBFGSStatus.INIT_FAILED_LOW_UPDATE_PCT: "INIT_FAILED_LOW_UPDATE_PCT: LBFGS failed to initialize due to the either: (1) LBFGS function or gradient values containing too many inf or nan values or (2) gradient changes being significantly large, set by epsilon. Consider reparameterizing the model, adjusting initvals or jitter or other pathfinder arguments if this occurence is high relative to the number of paths.",
    }

    path_status_message = {
        PathStatus.ELBO_ARGMAX_AT_ZERO: "ELBO_ARGMAX_AT_ZERO: ELBO argmax at zero refers to the first iteration during LBFGS. A high occurrence suggests the model's default initial point + jitter values are concentrated in high-density regions in the target distribution and may result in poor exploration of the parameter space. Consider increasing jitter if this occurrence is high relative to the number of paths.",
        PathStatus.INVALID_LOGQ: "INVALID_LOGQ: Invalid logQ values occur when a path's logQ values are not finite. The failed path is not included in samples when importance sampling is used. Consider reparameterizing the model or adjusting the pathfinder arguments if this occurence is high relative to the number of paths.",
        PathStatus.SINGLE_STEP: "SINGLE_STEP: Pathfinder requires at least two LBFGS steps on a path. A path with only one step produces an invalid result. Consider adjusting initvals/jitter if this occurs.",
    }

    for lbfgs_status in mpr.lbfgs_status:
        if lbfgs_status in lbfgs_status_message:
            warnings.append(lbfgs_status_message.get(lbfgs_status))

    for path_status in mpr.path_status:
        if path_status in path_status_message:
            warnings.append(path_status_message.get(path_status))

    return warnings


def _make_multipath_progress(progressbar: bool) -> CustomProgress:
    return CustomProgress(
        TextColumn("{task.description}", table_column=Column("Path", min_width=7, no_wrap=True)),
        TextColumn(
            "{task.fields[status]}", table_column=Column("Status", min_width=10, no_wrap=True)
        ),
        TextColumn(
            "{task.fields[lbfgs_steps]}", table_column=Column("Steps", min_width=6, no_wrap=True)
        ),
        TextColumn(
            "{task.fields[steps_per_sec]}",
            table_column=Column("Steps/s", min_width=8, no_wrap=True),
        ),
        TextColumn(
            "{task.fields[best_ind]}", table_column=Column("Best step", min_width=9, no_wrap=True)
        ),
        TextColumn(
            "{task.fields[best_elbo]}", table_column=Column("Best ELBO", min_width=12, no_wrap=True)
        ),
        TextColumn(
            "{task.fields[current_elbo]}",
            table_column=Column("Cur ELBO", min_width=12, no_wrap=True),
        ),
        TextColumn(
            "{task.fields[step_size]}",
            table_column=Column("Step size", min_width=10, no_wrap=True),
        ),
        TimeElapsedColumn(table_column=Column("Elapsed", min_width=8, no_wrap=True)),
        include_headers=True,
        console=Console(theme=default_progress_theme),
        disable=not progressbar,
    )


def _make_progress_callback(progress: CustomProgress, task_id: int) -> Callable[[dict], None]:
    def cb(info: dict) -> None:
        fields: dict[str, Any] = {}
        if "status" in info and info["status"] is not None:
            fields["status"] = info["status"]
        if "lbfgs_steps" in info:
            fields["lbfgs_steps"] = info["lbfgs_steps"]
        if "best_elbo" in info:
            val = info["best_elbo"]
            fields["best_elbo"] = (
                f"{val:.3f}" if val is not None and np.isfinite(float(val)) else "—"
            )
        if "best_ind" in info:
            val = info["best_ind"]
            fields["best_ind"] = (
                str(int(val)) if val is not None and np.isfinite(float(val)) else "—"
            )
        if "current_elbo" in info:
            val = info["current_elbo"]
            fields["current_elbo"] = (
                f"{val:.3f}" if val is not None and np.isfinite(float(val)) else "—"
            )
        if "step_size" in info:
            val = info["step_size"]
            fields["step_size"] = (
                f"{val:.2e}" if val is not None and np.isfinite(float(val)) else "—"
            )
        if "steps_per_sec" in info:
            val = info["steps_per_sec"]
            fields["steps_per_sec"] = (
                f"{val:.1f}/s" if val is not None and np.isfinite(float(val)) else "—"
            )
        if fields:
            progress.update(task_id, **fields)
        if info.get("status") in ("ok", "elbo@0"):
            progress.stop_task(task_id)

    return cb


def multipath_pathfinder(
    model: Model,
    num_paths: int,
    num_draws: int,
    num_draws_per_path: int,
    maxcor: int,
    maxiter: int,
    ftol: float,
    gtol: float,
    maxls: int,
    num_elbo_draws: int,
    jitter: float,
    epsilon: float,
    importance_sampling: Literal["psis", "psir", "identity"] | None,
    progressbar: bool,
    concurrent: Literal["process"] | None = "process",
    cores: int | None = None,
    blas_cores: int | None | Literal["auto"] = "auto",
    mp_ctx: mp.context.BaseContext | str | None = None,
    random_seed: RandomSeed = None,
    max_init_retries: int = 10,
    pathfinder_kwargs: dict[str, Any] = {},
    compile_kwargs: dict[str, Any] = {},
    display_summary: bool = True,
) -> MultiPathfinderResult:
    """
    Fit the Pathfinder Variational Inference algorithm using multiple paths with PyMC/PyTensor backend.

    Parameters
    ----------
    model : pymc.Model
        The PyMC model to fit the Pathfinder algorithm to.
    num_paths : int
        Number of independent paths to run in the Pathfinder algorithm. (default is 4) It is recommended to increase num_paths when increasing the jitter value.
    num_draws : int, optional
        Total number of samples to draw from the fitted approximation (default is 1000).
    num_draws_per_path : int, optional
        Number of samples to draw per path (default is 1000).
    maxcor : int, optional
        Maximum number of variable metric corrections used to define the limited memory matrix (default is None). If None, maxcor is set to ceil(3 * log(N)) or 5 whichever is greater, where N is the number of model parameters.
    maxiter : int, optional
        Maximum number of iterations for the L-BFGS optimisation (default is 1000).
    ftol : float, optional
        Tolerance for the decrease in the objective function (default is 1e-5).
    gtol : float, optional
        Tolerance for the norm of the gradient (default is 1e-8).
    maxls : int, optional
        Maximum number of line search steps for the L-BFGS algorithm (default is 1000).
    num_elbo_draws : int, optional
        Number of draws for the Evidence Lower Bound (ELBO) estimation (default is 10).
    jitter : float, optional
        Amount of jitter to apply to initial points (default is 2.0). Note that Pathfinder may be highly sensitive to the jitter value. It is recommended to increase num_paths when increasing the jitter value.
    epsilon: float
        value used to filter out large changes in the direction of the update gradient at each iteration l in L. Iteration l is only accepted if s·z >= epsilon * ||z||² for each l in L. Matches Zhang et al. (2022) Algorithm 3 with default epsilon=1e-12.
    importance_sampling : str, None, optional
        Method to apply sampling based on log importance weights (logP - logQ).
        "psis" : Pareto Smoothed Importance Sampling (default)
                Recommended for more stable results.
        "psir" : Pareto Smoothed Importance Resampling
                Less stable than PSIS.
        "identity" : Applies log importance weights directly without resampling.
        None : No importance sampling weights. Returns raw samples of size (num_paths, num_draws_per_path, N) where N is number of model parameters. Other methods return samples of size (num_draws, N).
    progressbar : bool, optional
        Whether to display a progress bar (default is False). Setting this to True will likely increase the computation time.
    random_seed : RandomSeed, optional
        Random seed for reproducibility.
    postprocessing_backend : str, optional
        Backend for postprocessing transformations, either "cpu" or "gpu" (default is "cpu"). This is only relevant if inference_backend is "blackjax".
    inference_backend : str, optional
        Backend for inference, either "pymc" or "blackjax" (default is "pymc").
    concurrent : str, optional
        How to run paths: ``"process"`` (default) spawns separate worker processes for true
        parallelism, matching PyMC's approach for parallel chains.  Set to ``None`` for serial
        execution (useful for debugging).
    cores : int, optional
        Number of paths to run in parallel. If ``None``, set to min(4, cpu_count, num_paths),
        mirroring pm.sample.
    blas_cores : int or "auto" or None, optional
        Total number of threads BLAS/OpenMP should use per worker. ``"auto"`` (default)
        matches total to ``cores``. ``None`` keeps default BLAS behavior.
    max_init_retries : int, optional
        Maximum number of re-jitter retries per path when LBFGSInitFailed is raised (default is 10).
    pathfinder_kwargs
        Additional keyword arguments for the Pathfinder algorithm. Supported keys:
        ``jacobian`` (bool, default True), ``vectorize`` (bool, default False).
    compile_kwargs
        Additional keyword arguments for the PyTensor compiler. If not provided, a
        performant default is used.

    Returns
    -------
    MultiPathfinderResult
        The result containing samples and other information from the Multi-Path Pathfinder algorithm.
    """

    *path_seeds, choice_seed = _get_seeds_per_chain(random_seed, num_paths + 1)

    pathfinder_config = PathfinderConfig(
        num_draws=num_draws_per_path,
        maxcor=maxcor,
        maxiter=maxiter,
        ftol=ftol,
        gtol=gtol,
        maxls=maxls,
        num_elbo_draws=num_elbo_draws,
        jitter=jitter,
        epsilon=epsilon,
    )

    compile_start = time.time()
    single_pathfinder_fn = make_single_pathfinder_fn(
        model,
        **asdict(pathfinder_config),
        max_init_retries=max_init_retries,
        pathfinder_kwargs=pathfinder_kwargs,
        compile_kwargs=compile_kwargs,
    )
    compile_end = time.time()

    results = []
    compute_start = time.time()
    try:
        # Per-path progress bar (one row per path, updated in real time)
        progress = _make_multipath_progress(progressbar)

        # Create one task per path and build per-path progress callbacks
        task_ids = []
        path_callbacks: list[Callable | None] = []
        with progress:
            for i in range(num_paths):
                tid = progress.add_task(
                    f"Path {i + 1}",
                    status="queued",
                    lbfgs_steps=0,
                    steps_per_sec="—",
                    best_elbo="—",
                    best_ind="—",
                    current_elbo="—",
                    step_size="—",
                    total=None,
                )
                task_ids.append(tid)
                path_callbacks.append(_make_progress_callback(progress, tid))

            # concurrent="process" gives true parallelism via separate worker processes
            # (matching PyMC's approach). concurrent=None is serial (useful for debugging).
            generator = make_generator(
                concurrent=concurrent,
                fn=single_pathfinder_fn,
                seeds=path_seeds,
                cores=cores,
                blas_cores=blas_cores,
                progress_callbacks=path_callbacks,
                mp_ctx=mp_ctx,
            )

            for result in generator:
                try:
                    if isinstance(result, Exception):
                        raise result
                    else:
                        results.append(result)
                except Exception as e:
                    logger.warning("Unexpected error in a path: %s", str(e))
                    results.append(
                        PathfinderResult(
                            path_status=PathStatus.PATH_FAILED,
                            lbfgs_status=LBFGSStatus.LBFGS_FAILED,
                        )
                    )
    except (KeyboardInterrupt, StopIteration) as e:
        # User is free to abort early — MultiPathfinderResult collects all successful results so far.
        if isinstance(e, StopIteration):
            logger.info(str(e))
    finally:
        compute_end = time.time()
        if results:
            mpr = (
                MultiPathfinderResult.from_path_results(results)
                .with_pathfinder_config(config=pathfinder_config)
                .with_importance_sampling(
                    num_draws=num_draws, method=importance_sampling, random_seed=choice_seed
                )
                .with_timing(
                    compile_time=compile_end - compile_start,
                    compute_time=compute_end - compute_start,
                )
            )
            # Display summary conditionally
            if display_summary:
                mpr.display_summary()
            if mpr.all_paths_failed:
                raise ValueError(
                    "All paths failed. Consider decreasing the jitter or reparameterizing the model."
                )
        else:
            raise ValueError(
                "BUG: Failed to iterate!"
                "Please report this issue at: "
                "https://github.com/pymc-devs/pymc-extras/issues "
                "with your code to reproduce the issue and the following details:\n"
                f"pathfinder_config: \n{pathfinder_config}\n"
                f"compile_kwargs: {compile_kwargs}\n"
                f"pathfinder_kwargs: {pathfinder_kwargs}\n"
                f"num_paths: {num_paths}\n"
                f"num_draws: {num_draws}\n"
            )

    return mpr


def fit_pathfinder(
    model=None,
    num_paths: int = 4,  # I
    num_draws: int = 1000,  # R
    num_draws_per_path: int = 1000,  # M
    maxcor: int | None = None,  # J
    maxiter: int = 1000,  # L^max
    ftol: float = 1e-5,
    gtol: float = 1e-8,
    maxls: int = 1000,
    num_elbo_draws: int = 10,  # K
    max_init_retries: int = 10,
    jitter: float = 2.0,
    epsilon: float = 1e-12,
    importance_sampling: Literal["psis", "psir", "identity"] | None = "psis",
    progressbar: bool = True,
    concurrent: Literal["process"] | None = "process",
    cores: int | None = None,
    blas_cores: int | None | Literal["auto"] = "auto",
    mp_ctx: mp.context.BaseContext | str | None = None,
    random_seed: RandomSeed | None = None,
    postprocessing_backend: Literal["cpu", "gpu"] = "cpu",
    inference_backend: Literal["pymc", "blackjax"] = "pymc",
    pathfinder_kwargs: dict[str, Any] = {},
    compile_kwargs: dict[str, Any] = {},
    initvals: dict[str, Any] | None = None,
    # New pathfinder result integration options
    add_pathfinder_groups: bool = True,
    display_summary: bool | Literal["auto"] = "auto",
    store_diagnostics: bool = False,
) -> az.InferenceData:
    """
    Fit the Pathfinder Variational Inference algorithm.

    This function fits the Pathfinder algorithm to a given PyMC model, allowing for multiple paths and draws. It supports both PyMC and BlackJAX backends.

    Parameters
    ----------
    model : pymc.Model
        The PyMC model to fit the Pathfinder algorithm to.
    num_paths : int
        Number of independent paths to run in the Pathfinder algorithm. (default is 4) It is recommended to increase num_paths when increasing the jitter value.
    num_draws : int, optional
        Total number of samples to draw from the fitted approximation (default is 1000).
    num_draws_per_path : int, optional
        Number of samples to draw per path (default is 1000).
    maxcor : int, optional
        Maximum number of variable metric corrections used to define the limited memory matrix (default is None). If None, maxcor is set to ceil(3 * log(N)) or 5 whichever is greater, where N is the number of model parameters.
    maxiter : int, optional
        Maximum number of iterations for the L-BFGS optimisation (default is 1000).
    ftol : float, optional
        Tolerance for the decrease in the objective function (default is 1e-5).
    gtol : float, optional
        Tolerance for the norm of the gradient (default is 1e-8).
    maxls : int, optional
        Maximum number of line search steps for the L-BFGS algorithm (default is 1000).
    num_elbo_draws : int, optional
        Number of draws for the Evidence Lower Bound (ELBO) estimation (default is 10).
    jitter : float, optional
        Amount of jitter to apply to initial points (default is 2.0). Note that Pathfinder may be highly sensitive to the jitter value. It is recommended to increase num_paths when increasing the jitter value.
    epsilon: float
        value used to filter out large changes in the direction of the update gradient at each iteration l in L. Iteration l is only accepted if s·z >= epsilon * ||z||² for each l in L. Matches Zhang et al. (2022) Algorithm 3 with default epsilon=1e-12.
    importance_sampling : str, None, optional
        Method to apply sampling based on log importance weights (logP - logQ).
        Options are:

        - "psis" : Pareto Smoothed Importance Sampling (default). Usually most stable.
        - "psir" : Pareto Smoothed Importance Resampling. Less stable than PSIS.
        - "identity" : Applies log importance weights directly without resampling.
        - None : No importance sampling weights. Returns raw samples of size (num_paths, num_draws_per_path, N) where N is number of model parameters. Other methods return samples of size (num_draws, N).

    progressbar : bool, optional
        Whether to display a progress bar (default is True). Setting this to False will likely reduce the computation time.
    random_seed : RandomSeed, optional
        Random seed for reproducibility.
    postprocessing_backend : str, optional
        Backend for postprocessing transformations, either "cpu" or "gpu" (default is "cpu"). This is only relevant if inference_backend is "blackjax".
    inference_backend : str, optional
        Backend for inference, either "pymc" or "blackjax" (default is "pymc").
    concurrent : str, optional
        How to run paths: ``"process"`` (default) spawns separate worker processes for true
        parallelism, matching PyMC's approach for parallel chains.  Set to ``None`` for serial
        execution (useful for debugging).
    cores : int, optional
        Number of paths to run in parallel. If ``None``, set to min(4, cpu_count, num_paths),
        mirroring pm.sample.
    blas_cores : int or "auto" or None, optional
        Total number of threads BLAS/OpenMP should use per worker. ``"auto"`` (default)
        matches total to ``cores``. ``None`` keeps default BLAS behavior.
    mp_ctx : str or multiprocessing.Context, optional
        Multiprocessing context for parallel path execution (e.g. ``"spawn"``, ``"fork"``).
    max_init_retries : int, optional
        Maximum number of re-jitter retries per path when the initial point fails (default is 10).
    pathfinder_kwargs : dict, optional
        Additional keyword arguments for the Pathfinder algorithm. Supported keys:
        - ``vectorize`` (bool, default False): If True, use ``vectorize_graph``
          instead of ``pytensor.map`` for batched log-probability evaluation.
          May be faster for small models but may use more memory for large
          parameter counts.
        - ``jacobian`` (bool, default True): Whether to include the Jacobian
          determinant in the log-probability computation. Setting to False is not
          recommended and may result in very high Pareto k values.

    compile_kwargs
        Additional keyword arguments for the PyTensor compiler. If not provided, a
        performant default is used.
    initvals: dict | None = None
        Initial values for the model parameters, as str:ndarray key-value pairs. Paritial initialization is permitted.
        If None, the model's default initial values are used.
    add_pathfinder_groups : bool, optional
        Whether to add pathfinder results as additional groups to the InferenceData (default is True).
        When True, adds pathfinder and pathfinder_paths groups with optimization diagnostics.
    display_summary : bool or "auto", optional
        Whether to display the pathfinder results summary (default is "auto").
        "auto" preserves current behavior, False suppresses console output.
    store_diagnostics : bool, optional
        Whether to include potentially large diagnostic arrays in the pathfinder groups (default is False).
    pathfinder_group : str, optional
        Name for the main pathfinder results group (default is "sample_stats").

    Returns
    -------
    :class:`~arviz.InferenceData`
        The inference data containing the results of the Pathfinder algorithm.

    References
    ----------
    Zhang, L., Carpenter, B., Gelman, A., & Vehtari, A. (2022). Pathfinder: Parallel quasi-Newton variational inference. Journal of Machine Learning Research, 23(306), 1-49.
    """

    model = modelcontext(model)
    if concurrent is not None and concurrent != "process":
        raise ValueError(
            f"concurrent must be 'process' or None, got {concurrent!r}. "
            "Thread-based parallelism is not supported: PyTensor compiled functions "
            "share internal storage and would corrupt each other's state."
        )

    if initvals is not None:
        model = pm.model.fgraph.clone_model(model)  # Create a clone of the model
        for (
            rv_name,
            ivals,
        ) in initvals.items():  # Set the initial values for the variables in the clone
            model.set_initval(model.named_vars[rv_name], ivals)

    valid_importance_sampling = {"psis", "psir", "identity", None}

    if importance_sampling is not None:
        importance_sampling = importance_sampling.lower()

    if importance_sampling not in valid_importance_sampling:
        raise ValueError(f"Invalid importance sampling method: {importance_sampling}")

    N = DictToArrayBijection.map(model.initial_point()).data.shape[0]

    if maxcor is None:
        # Based on tests, this seems to be a good default value. Higher maxcor values do not necessarily lead to better results and can slow down the algorithm. Also, if results do benefit from a higher maxcor value, the improvement may be diminishing w.r.t. the increase in maxcor.
        maxcor = np.ceil(3 * np.log(N)).astype(np.int32)
        maxcor = max(maxcor, 5)

    # Handle display_summary logic
    should_display_summary = display_summary == "auto" or display_summary is True

    if inference_backend == "pymc":
        mp_result = multipath_pathfinder(
            model,
            num_paths=num_paths,
            num_draws=num_draws,
            num_draws_per_path=num_draws_per_path,
            maxcor=maxcor,
            maxiter=maxiter,
            ftol=ftol,
            gtol=gtol,
            maxls=maxls,
            num_elbo_draws=num_elbo_draws,
            jitter=jitter,
            epsilon=epsilon,
            importance_sampling=importance_sampling,
            progressbar=progressbar,
            concurrent=concurrent,
            cores=cores,
            blas_cores=blas_cores,
            mp_ctx=mp_ctx,
            max_init_retries=max_init_retries,
            random_seed=random_seed,
            pathfinder_kwargs=pathfinder_kwargs,
            compile_kwargs=compile_kwargs,
            display_summary=should_display_summary,
        )
        pathfinder_samples = mp_result.samples
    elif inference_backend == "blackjax":
        import blackjax
        import jax

        if version.parse(blackjax.__version__).major < 1:
            raise ImportError("fit_pathfinder requires blackjax 1.0 or above")

        _, pathfinder_seed, sample_seed = _get_seeds_per_chain(random_seed, 3)
        # TODO: extend initial points with jitter_scale to blackjax
        # TODO: extend blackjax pathfinder to multiple paths
        x0, _ = DictToArrayBijection.map(model.initial_point())
        logp_func = get_jaxified_logp_of_ravel_inputs(model)
        pathfinder_state, _ = blackjax.vi.pathfinder.approximate(
            rng_key=jax.random.key(pathfinder_seed),
            logdensity_fn=logp_func,
            initial_position=x0,
            num_samples=num_elbo_draws,
            maxiter=maxiter,
            maxcor=maxcor,
            maxls=maxls,
            ftol=ftol,
            gtol=gtol,
            **pathfinder_kwargs,
        )
        pathfinder_samples, _ = blackjax.vi.pathfinder.sample(
            rng_key=jax.random.key(sample_seed),
            state=pathfinder_state,
            num_samples=num_draws,
        )
    else:
        raise ValueError(f"Invalid inference_backend: {inference_backend}")

    logger.info("Transforming variables...")

    idata = convert_flat_trace_to_idata(
        pathfinder_samples,
        postprocessing_backend=postprocessing_backend,
        inference_backend=inference_backend,
        model=model,
        importance_sampling=importance_sampling,
    )

    idata = add_data_to_inference_data(idata, progressbar, model, compile_kwargs)

    # Add pathfinder results to InferenceData if requested
    if add_pathfinder_groups:
        if inference_backend == "pymc":
            from pymc_extras.inference.pathfinder.idata import add_pathfinder_to_inference_data

            idata = add_pathfinder_to_inference_data(
                idata=idata,
                result=mp_result,
                model=model,
                store_diagnostics=store_diagnostics,
            )
        else:
            warnings.warn(
                f"Pathfinder diagnostic groups are only supported with the PyMC backend. "
                f"Current backend is '{inference_backend}', which does not support adding "
                "pathfinder diagnostics to InferenceData. The InferenceData will only contain "
                "posterior samples. To add diagnostic groups, use inference_backend='pymc', "
                "or set add_pathfinder_groups=False to suppress this warning.",
                UserWarning,
                stacklevel=2,
            )

    return idata
