code: def lbfgs(
      *,
      log_target_density: Callable[[jnp.ndarray], jnp.ndarray],
      initial_value: jnp.ndarray,  # theta_init
      inverse_hessian_history: int = 6,  # J
      relative_tolerance: float = 1e-13,  # tau_rel
      max_iters: int = 1000,  # L
      wolfe_bounds: Tuple[float, float] = (1e-4, 0.9),
      positivity_threshold: float = 2.2e-16,
  ):
      """LBFGS implementation which returns the optimization path and gradients."""
      dim = initial_value.shape[0]
      grad_log_density = jax.grad(log_target_density)
      optimization_path = [initial_value]
      current_lp = log_target_density(initial_value)
      grad_optimization_path = [grad_log_density(initial_value)]
      position_diffs = jnp.empty((dim, 0))
      gradient_diffs = jnp.empty((dim, 0))
      for _ in range(max_iters):
          (
              diagonal_estimate,
              thin_factors,
              scaling_outer_product,
          ) = bfgs_inverse_hessian(
              updates_of_position_differences=position_diffs,
              updates_of_gradient_differences=gradient_diffs,
          )
          grad_lp = grad_optimization_path[-1]
          search_direction = (
              diagonal_estimate * grad_lp
              + thin_factors
              @ (
                  scaling_outer_product
                  @ (jnp.transpose(thin_factors) @ grad_lp)
              )
          )
          step_size = 1.0
          while step_size > 1e-8:
              proposed = (
                  optimization_path[-1] + step_size * search_direction
              )
              proposed_lp = log_target_density(proposed)
              if proposed_lp >= current_lp + (
                  wolfe_bounds[0] * grad_lp
              ) @ (step_size * search_direction):
                  proposed_grad = grad_log_density(proposed)
                  if (
                      proposed_grad @ search_direction
                      <= wolfe_bounds[1] * grad_lp @ search_direction
                  ):
                      break
              step_size = 0.5 * step_size
          optimization_path.append(proposed)
          grad_optimization_path.append(proposed_grad)
          if (proposed_lp - current_lp) / jnp.abs(
              current_lp
          ) < relative_tolerance:
              return optimization_path, grad_optimization_path
          current_lp = proposed_lp
  
          position_diff: jnp.ndarray = (
              optimization_path[-1] - optimization_path[-2]
          )
          grad_diff = (
              -grad_optimization_path[-1] + grad_optimization_path[-2]
          )
          if (
              position_diff @ grad_diff
              > positivity_threshold * jnp.sum(grad_diff**2)
          ):
              position_diffs = jnp.column_stack(
                  (
                      position_diffs[:, -inverse_hessian_history + 1 :],
                      position_diff,
                  )
              )
              gradient_diffs = jnp.column_stack(
                  (
                      gradient_diffs[:, -inverse_hessian_history + 1 :],
                      grad_diff,
                  )
              )
      return optimization_path, grad_optimization_path
  
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 6 for calling a global function with example usage: lbfgs() and returns: optimization_path,  grad_optimization_path, optimization_path,  grad_optimization_path