code: def _update_inv_hessian(prev_state, next_state):
      """Update the BGFS state by computing the next inverse hessian estimate."""
      # Only update the inverse Hessian if not already failed or converged.
      should_update = ~next_state.converged & ~next_state.failed
  
      # Compute the normalization term (y^T . s), should not update if is singular.
      gradient_delta = (
          next_state.objective_gradient - prev_state.objective_gradient
      )
      position_delta = next_state.position - prev_state.position
      normalization_factor = tf.reduce_sum(
          gradient_delta * position_delta, axis=-1
      )
      should_update = should_update & ~tf.equal(normalization_factor, 0)
  
      # Rescale the initial hessian at the first step, as suggested
      # in Chapter 6 of Numerical Optimization, by Nocedal and Wright.
      scale_factor = tf.where(
          (
              tf.math.equal(prev_state.num_iterations, 0)
              & prev_state.scale_initial_inverse_hessian
          ),
          normalization_factor
          / tf.reduce_sum(tf.math.square(gradient_delta), axis=-1),
          1.0,
      )
  
      inverse_hessian_estimate = (
          scale_factor[..., tf.newaxis, tf.newaxis]
          * prev_state.inverse_hessian_estimate
      )
  
      def _do_update_inv_hessian():
          next_inv_hessian = _bfgs_inv_hessian_update(
              gradient_delta,
              position_delta,
              normalization_factor,
              inverse_hessian_estimate,
          )
          return bfgs_utils.update_fields(
              next_state,
              inverse_hessian_estimate=tf.where(
                  should_update[..., tf.newaxis, tf.newaxis],
                  next_inv_hessian,
                  inverse_hessian_estimate,
              ),
          )
  
      return ps.cond(
          ps.reduce_any(should_update),
          _do_update_inv_hessian,
          lambda: next_state,
      )
  
file_path: probability/tensorflow_probability/python/optimizer/bfgs.py
imports: import collections
  import tensorflow.compat.v2
  from tensorflow_probability.python.internal import distribution_util
  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 4 for calling a global function with example usage: _update_inv_hessian(prev_state, next_state) and returns: ps, bfgs_utils