code: def line_search_step(
      state,
      value_and_gradients_function,
      search_direction,
      grad_tolerance,
      f_relative_tolerance,
      x_tolerance,
      stopping_condition,
      max_iterations,
      f_absolute_tolerance,
  ):
      """Performs the line search step of the BFGS search procedure.
  
      Uses hager_zhang line search procedure to compute a suitable step size
      to advance the current `state.position` along the given `search_direction`.
      Also, if the line search is successful, updates the `state.position` by
      taking the corresponding step.
  
      Args:
        state: A namedtuple instance holding values for the current state of the
          search procedure. The state must include the fields: `position`,
            `objective_value`, `objective_gradient`, `num_iterations`,
            `num_objective_evaluations`, `converged` and `failed`.
        value_and_gradients_function: A Python callable that accepts a point as a
          real `Tensor` of shape `[..., n]` and returns a tuple of two tensors of
          the same dtype: the objective function value, a real `Tensor` of shape
            `[...]`, and its derivative, another real `Tensor` of shape `[..., n]`.
        search_direction: A real `Tensor` of shape `[..., n]`. The direction along
          which to perform line search.
        grad_tolerance: Scalar `Tensor` of real dtype. Specifies the gradient
          tolerance for the procedure.
        f_relative_tolerance: Scalar `Tensor` of real dtype. Specifies the tolerance
          for the relative change in the objective value.
        x_tolerance: Scalar `Tensor` of real dtype. Specifies the tolerance for the
          change in the position.
        stopping_condition: A Python function that takes as input two Boolean
          tensors of shape `[...]`, and returns a Boolean scalar tensor. The input
          tensors are `converged` and `failed`, indicating the current status of
          each respective batch member; the return value states whether the
          algorithm should stop.
        max_iterations: A Python integer that is used as the maximum number of
          iterations of the hager_zhang line search algorithm
        f_absolute_tolerance: Scalar `Tensor` of real dtype. Specifies the tolerance
          for the absolute change in the objective value.
  
      Returns:
        A copy of the input state with the following fields updated:
          converged: a Boolean `Tensor` of shape `[...]` indicating whether the
            convergence criteria has been met.
          failed: a Boolean `Tensor` of shape `[...]` indicating whether the line
            search procedure failed to converge, or if either the updated gradient
            or objective function are no longer finite.
          num_iterations: Increased by 1.
          num_objective_evaluations: Increased by the number of times that the
            objective function got evaluated.
          position, objective_value, objective_gradient: If line search succeeded,
            updated by computing the new position and evaluating the objective
            function at that position.
      """
      line_search_value_grad_func = _restrict_along_direction(
          value_and_gradients_function, state.position, search_direction
      )
      derivative_at_start_pt = tf.reduce_sum(
          state.objective_gradient * search_direction, axis=-1
      )
      val_0 = ValueAndGradient(
          x=_broadcast(0, state.position),
          f=state.objective_value,
          df=derivative_at_start_pt,
          full_gradient=state.objective_gradient,
      )
      inactive = state.failed | state.converged
      ls_result = hager_zhang(
          line_search_value_grad_func,
          initial_step_size=_broadcast(1, state.position),
          value_at_zero=val_0,
          converged=inactive,
          max_iterations=max_iterations,
      )  # No search needed for these.
  
      state_after_ls = update_fields(
          state,
          failed=state.failed
          | (~state.converged & ~ls_result.converged),
          num_iterations=state.num_iterations + 1,
          num_objective_evaluations=(
              state.num_objective_evaluations + ls_result.func_evals
          ),
      )
  
      def _do_update_position():
          # For inactive batch members `left.x` is zero. However, their
          # `search_direction` might also be undefined, so we can't rely on
          # multiplication by zero to produce a `position_delta` of zero.
          position_delta = tf.where(
              inactive[..., tf.newaxis],
              dtype_util.as_numpy_dtype(search_direction.dtype)(0),
              search_direction * ls_result.left.x[..., tf.newaxis],
          )
          return _update_position(
              state_after_ls,
              position_delta,
              ls_result.left.f,
              ls_result.left.full_gradient,
              grad_tolerance,
              f_relative_tolerance,
              x_tolerance,
              f_absolute_tolerance,
          )  # pyformat: disable
  
      return ps.cond(
          stopping_condition(state.converged, state.failed),
          true_fn=lambda: state_after_ls,
          false_fn=_do_update_position,
      )
  
file_path: probability/tensorflow_probability/python/optimizer/bfgs_utils.py
imports: import collections
  import numpy
  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.linesearch.hager_zhang import hager_zhang
label_desc: Use image node_id 6 for calling a global function with example usage: line_search_step(state, value_and_gradients_function, search_direction, grad_tolerance, f_relative_tolerance, x_tolerance, stopping_condition, max_iterations, f_absolute_tolerance) and returns: ps, _update_position