code: def _get_search_direction(state):
      """Computes the search direction to follow at the current state.
  
      On the `k`-th iteration of the main L-BFGS algorithm, the state has collected
      the most recent `m` correction pairs in position_deltas and gradient_deltas,
      where `k = state.num_iterations` and `m = min(k, num_correction_pairs)`.
  
      Assuming these, the code below is an implementation of the L-BFGS two-loop
      recursion algorithm given by [Nocedal and Wright(2006)][1]:
  
      ```None
        q_direction = objective_gradient
        for i in reversed(range(m)):  # First loop.
          inv_rho[i] = gradient_deltas[i]^T * position_deltas[i]
          alpha[i] = position_deltas[i]^T * q_direction / inv_rho[i]
          q_direction = q_direction - alpha[i] * gradient_deltas[i]
  
        kth_inv_hessian_factor = (gradient_deltas[-1]^T * position_deltas[-1] /
                                  gradient_deltas[-1]^T * gradient_deltas[-1])
        r_direction = kth_inv_hessian_factor * I * q_direction
  
        for i in range(m):  # Second loop.
          beta = gradient_deltas[i]^T * r_direction / inv_rho[i]
          r_direction = r_direction + position_deltas[i] * (alpha[i] - beta)
  
        return -r_direction  # Approximates - H_k * objective_gradient.
      ```
  
      Args:
        state: A `LBfgsOptimizerResults` tuple with the current state of the
          search procedure.
  
      Returns:
        A real `Tensor` of the same shape as the `state.position`. The direction
        along which to perform line search.
      """
      # The number of correction pairs that have been collected so far.
      num_elements = ps.minimum(
          state.num_iterations,  # TODO(b/162733947): Change loop state -> closure.
          ps.shape(state.position_deltas)[0],
      )
  
      def _two_loop_algorithm():
          """L-BFGS two-loop algorithm."""
          # Correction pairs are always appended to the end, so only the latest
          # `num_elements` vectors have valid position/gradient deltas. Vectors
          # that haven't been computed yet are zero.
          position_deltas = state.position_deltas
          gradient_deltas = state.gradient_deltas
  
          # Pre-compute all `inv_rho[i]`s.
          inv_rhos = tf.reduce_sum(
              gradient_deltas * position_deltas, axis=-1
          )
  
          def first_loop(acc, args):
              _, q_direction = acc
              position_delta, gradient_delta, inv_rho = args
              alpha = tf.math.divide_no_nan(
                  tf.reduce_sum(position_delta * q_direction, axis=-1),
                  inv_rho,
              )
              direction_delta = alpha[..., tf.newaxis] * gradient_delta
              return (alpha, q_direction - direction_delta)
  
          # Run first loop body computing and collecting `alpha[i]`s, while also
          # computing the updated `q_direction` at each step.
          zero = tf.zeros_like(inv_rhos[-num_elements])
          alphas, q_directions = tf.scan(
              first_loop,
              [position_deltas, gradient_deltas, inv_rhos],
              initializer=(zero, state.objective_gradient),
              reverse=True,
          )
  
          # We use `H^0_k = gamma_k * I` as an estimate for the initial inverse
          # hessian for the k-th iteration; then `r_direction = H^0_k * q_direction`.
          gamma_k = inv_rhos[-1] / tf.reduce_sum(
              gradient_deltas[-1] * gradient_deltas[-1], axis=-1
          )
          r_direction = (
              gamma_k[..., tf.newaxis] * q_directions[-num_elements]
          )
  
          def second_loop(r_direction, args):
              alpha, position_delta, gradient_delta, inv_rho = args
              beta = tf.math.divide_no_nan(
                  tf.reduce_sum(gradient_delta * r_direction, axis=-1),
                  inv_rho,
              )
              direction_delta = (alpha - beta)[
                  ..., tf.newaxis
              ] * position_delta
              return r_direction + direction_delta
  
          # Finally, run second loop body computing the updated `r_direction` at each
          # step.
          r_directions = tf.scan(
              second_loop,
              [alphas, position_deltas, gradient_deltas, inv_rhos],
              initializer=r_direction,
          )
          return -r_directions[-1]
  
      return ps.cond(
          ps.equal(num_elements, 0),
          lambda: -state.objective_gradient,
          _two_loop_algorithm,
      )
  
file_path: probability/tensorflow_probability/python/optimizer/lbfgs.py
imports: import collections
  import tensorflow.compat.v2
  from tensorflow_probability.python.internal import dtype_util
  from tensorflow_probability.python.internal import prefer_static
  from tensorflow_probability.python.optimizer import bfgs_utils
label_desc: Use image node_id 3 for calling a global function with example usage: _get_search_direction(state) and returns: ps, unknown, alpha,  unknown, unknown