fit_pathfinder#

pymc_extras.inference.fit_pathfinder(model=None, num_paths: int = 4, num_draws: int = 1000, num_draws_per_path: int = 1000, num_elbo_draws: int = 10, max_init_retries: int = 10, jitter: float = 2.0, lbfgs_config: LBFGSConfig | None = None, importance_sampling: Literal['psis', 'psir', 'identity'] | None = 'psis', progressbar: bool = True, parallel: bool = True, cores: int | None = None, blas_cores: int | None | Literal['auto'] = 'auto', mp_ctx: BaseContext | str | None = None, random_seed: None | int | Sequence[int] | ndarray = None, jacobian_correction: bool = True, vectorize_logp: bool = True, vectorize_postprocessing: bool = True, compile_kwargs: dict[str, Any] | None = None, initvals: dict[str, Any] | None = None) → DataTree[source]#

Fit Pathfinder variational inference (multi-path, PyMC/PyTensor backend).

For the blackjax-backed single-path variant, see fit_blackjax_pathfinder().

Parameters:

  • model (pymc.Model) – The PyMC model to fit the Pathfinder algorithm to.

  • num_paths (int, optional) – Number of independent paths to run. Increase this when increasing the jitter value. Default 4.

  • num_draws (int, optional) – Total number of samples to draw from the fitted approximation. Default 1000.

  • num_draws_per_path (int, optional) – Number of samples to draw per path. Default 1000.

  • num_elbo_draws (int, optional) – Number of draws for the Evidence Lower Bound (ELBO) estimation. Default 10.

  • max_init_retries (int, optional) – Maximum number of re-jitter retries per path when the initial point fails. Default 10.

  • jitter (float, optional) – Amount of jitter to apply to initial points. Pathfinder can be highly sensitive to this value; increase num_paths when increasing it. Default 2.0.

  • lbfgs_config (LBFGSConfig, optional) – L-BFGS configuration. For details, including default arguments, see LBFGSConfig.

  • importance_sampling (str or None, optional) –

    Method to apply based on log importance weights (logP - logQ):

    • ”psis” : Pareto Smoothed Importance Sampling; usually most stable.

    • ”psir” : Pareto Smoothed Importance Resampling; less stable than PSIS.

    • ”identity” : apply log importance weights directly without resampling.

    • None : no importance sampling; return raw samples of shape (num_paths, num_draws_per_path, N). The other methods return shape (num_draws, N).

    Default “psis”.

  • progressbar (bool, optional) – Whether to display a progress bar. Disabling it likely reduces computation time. Default True.

  • parallel (bool, optional) – If True, spawn a separate worker process per path for true parallelism (matching PyMC’s approach for parallel chains). If False, run paths serially in the main process, which is useful for debugging. Default True.

  • cores (int, optional) – Number of paths to run in parallel. If None, set to min(4, cpu_count, num_paths), mirroring pm.sample. Default None.

  • blas_cores (int or "auto" or None, optional) – Total number of threads BLAS/OpenMP should use per worker. “auto” matches the total to cores; None keeps default BLAS behavior. Default “auto”.

  • mp_ctx (str or multiprocessing.Context, optional) – Multiprocessing context for parallel path execution (e.g. "spawn", "fork").

  • random_seed (RandomSeed, optional) – Random seed for reproducibility.

  • jacobian_correction (bool, optional) – Whether to add the log-determinant-of-Jacobian correction term to model.logp to account for value-var transforms (e.g. log, logit). With the correction, logp is the joint density on unconstrained coordinates, which is what L-BFGS optimizes and what importance sampling needs. Disabling it generally produces very high pareto-k values. Default True.

  • vectorize_logp (bool, optional) – If True, use vectorize_graph to batch model.logp across the num_draws axis for ELBO and importance-sampling evaluation; if False, fall back to pytensor.map. This trades high memory with parallel compute (True) against low memory with sequential compute (False); prefer True unless the model is memory bound. Default True.

  • vectorize_postprocessing (bool, optional) – If True, use vectorize_graph to batch the Deterministic post-processing subgraph across all draws in one call; if False, iterate draws with pytensor.scan. Set to False when memory is a concern, e.g. with large intermediate computations. Default True.

  • compile_kwargs (dict, optional) – Additional keyword arguments for the PyTensor compiler. Default None.

  • initvals (dict, optional) – Initial values for the model parameters, as name-to-ndarray pairs. Partial initialization is permitted. If None, the model’s default initial values are used. Default None.

Returns:

The inference data containing the results of the Pathfinder algorithm.

Return type:

DataTree

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.