code: def pathfinder(
      *,
      target_density: Callable[[jnp.ndarray], jnp.ndarray],
      initial_value: jnp.ndarray,
      lbfgs_max_iters: int,  # L
      key,
      lbfgs_relative_tolerance: float = 1e-13,  # tau_rel
      lbfgs_inverse_hessian_history: int = 6,  # J
      monte_carlo_elbo_draws: int = 10,  # K
      num_draws: int = 10,  # M
  ):
      """Single pathfinder algorithm."""
  
      optimization_path, grad_optimization_path = lbfgs(
          log_target_density=target_density,
          initial_value=initial_value,
          inverse_hessian_history=lbfgs_inverse_hessian_history,
          relative_tolerance=lbfgs_relative_tolerance,
          max_iters=lbfgs_max_iters,
      )
      cov_factors = cov_estimate(
          optimization_path=optimization_path,
          optimization_path_grads=grad_optimization_path,
          history=lbfgs_inverse_hessian_history,
      )
      key, *keys = jax.random.split(key, num=len(optimization_path))
      elbos = np.empty(len(optimization_path) - 1)
      for idx, (value, factors, sample_key) in enumerate(
          zip(optimization_path, cov_factors, keys)
      ):
          (
              diagonal_estimate,
              thin_factors,
              scaling_outer_product,
          ) = factors
          draws, draw_probs = bfgs_sample(
              value=value,
              grad_density=grad_optimization_path[idx],
              diagonal_estimate=diagonal_estimate,
              thin_factors=thin_factors,
              scaling_outer_product=scaling_outer_product,
              num_draws=monte_carlo_elbo_draws,
              key=sample_key,
          )
          target_probs = target_density(draws)
  
          elbos[idx] = elbo(target_probs, draw_probs)
      best_idx = jnp.argmax(elbos)
      (
          diagonal_estimate,
          thin_factors,
          scaling_outer_product,
      ) = cov_factors[best_idx]
      return bfgs_sample(
          value=optimization_path[best_idx],
          grad_density=grad_optimization_path[best_idx],
          diagonal_estimate=diagonal_estimate,
          thin_factors=thin_factors,
          scaling_outer_product=scaling_outer_product,
          num_draws=num_draws,
          key=key,
      )
  
file_path: probability/discussion/pathfinder/pathfinder.py
imports: from typing import Tuple, List, Sequence, Callable
  import jax
  import jax.numpy
  import numpy
label_desc: Use image node_id 1 for calling a global function with example usage: pathfinder() and returns: bfgs_sample