code: def bfgs_sample(
      *,
      value,
      grad_density,
      diagonal_estimate,
      thin_factors,
      scaling_outer_product,
      num_draws,
      key,
  ):
      """Sample from BFGS factorized estimates."""
      q, r = jnp.linalg.qr(
          jnp.diag(jax.lax.rsqrt(diagonal_estimate)) @ thin_factors
      )
      lower_chol = jnp.linalg.cholesky(
          jnp.eye(r.shape[0]) + r @ scaling_outer_product @ r.T
      )
      log_sigma = (
          jnp.log(jnp.prod(jnp.abs(diagonal_estimate)))
          + 2.0 * jnp.linalg.slogdet(lower_chol)[1]
      )
      mu = (
          value
          + diagonal_estimate * grad_density
          + thin_factors
          @ (
              scaling_outer_product
              @ (jnp.transpose(thin_factors) @ grad_density)
          )
      )
      draws = jax.random.multivariate_normal(
          key=key,
          mean=jnp.zeros_like(value),
          cov=jnp.eye(value.shape[0]),
          shape=(num_draws,),
      ).T
      transformed = q.T @ draws
      transformed_draws = (
          mu
          + (
              jnp.diag(jnp.sqrt(diagonal_estimate))
              @ (q @ lower_chol @ transformed + draws - q @ transformed)
          ).T
      )
      draw_probs = -0.5 * (
          log_sigma
          + jnp.sum(draws**2, axis=0)
          + value.shape[0] * jnp.log(2 * jnp.pi)
      )
  
      return transformed_draws, draw_probs
  
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 5 for calling a global function with example usage: bfgs_sample() and returns: transformed_draws,  draw_probs