class_docstr: Base class for stopping rules.
  
  A convergence criterion determines when an optimization has converged given
  its history of losses, gradients, and parameter values. Each criterion is
  responsible for propagating from step to step whatever state it needs to
  represent the relevant aspects of that history (for example, a moving average
  of previous loss values or gradients). In particular, subclasses must
  implement:
  
  - `_bootstrap(loss, grads, parameters)`: takes the
    initial loss, gradients, and values of parameters, and returns a (structure
    of) `Tensor`(s) representing the initial values of any auxiliary quantities
    tracked by the convergence criterion.
  - `_one_step(step, loss, grads, window_size, auxiliary_state)`: At
    integer `step >= 1`,
    takes the current loss, gradients, and values of parameters, along with
    any auxiliary state carried over from the previous step, and returns
    `(has_converged, updated_auxiliary_state)`, where `has_converged` is a
    boolean `Tensor`, and `updated_auxiliary_state` is a (structure of)
    Tensor(s) matching `auxiliary_state`, containing whatever information must
    be propagated to the next timestep.
docstr: Updates tracked quantities for a new step, and determines if converged.
  
  The shape of the `Tensor`s specifying `loss`, `grads`, and `parameters` may
  optionally be prefixed by one or more batch dimension(s). In this case,
  the returned value `has_converged` will have shape equal to the broadcast
  batch shape of whichever of those quantities is used by this convergence
  criterion, and the quantities defining the convergence criterion (
  `min_num_steps`, etc.).
  
  Args:
    step: integer `Tensor` index of the current step, where `step >= 1` (on
      step `0`, `initial_state` should be called instead).
    loss: float `Tensor`  value of loss at the current step.
    grads: list of float `Tensor` gradients of `loss` wrt `parameters`.
    parameters: list of float `Tensor` current values of parameters
      being optimized.
    auxiliary_state: the (structure of) `Tensor`(s) containing state carried
      forward from the previous step.
  Returns:
    has_converged: boolean `Tensor` indicating whether the optimization has
      converged.
    updated_auxiliary_state: (Structure of) `Tensor`(s) representing
      updated quantities tracked by the convergence criterion. This should
      match the structure of the value returned by `bootstrap`.
docstr_tok: ["Updates","tracked","quantities","for","a","new","step",",","and","determines","if","converged",".","The","shape","of","the","`","Tensor","`","s","specifying","`","loss","`",",","`","grads","`",",","and","`","parameters","`","may","optionally","be","prefixed","by","one","or","more","batch","dimension","(","s",")",".","In","this","case",",","the","returned","value","`","has_converged","`","will","have","shape","equal","to","the","broadcast","batch","shape","of","whichever","of","those","quantities","is","used","by","this","convergence","criterion",",","and","the","quantities","defining","the","convergence","criterion","(","`","min_num_steps","`",",","etc.",")",".","Args",":","step",":","integer","`","Tensor","`","index","of","the","current","step",",","where","`","step",">","=","1","`","(","on","step","`","0","`",",","`","initial_state","`","should","be","called","instead",")",".","loss",":","float","`","Tensor","`","value","of","loss","at","the","current","step",".","grads",":","list","of","float","`","Tensor","`","gradients","of","`","loss","`","wrt","`","parameters","`",".","parameters",":","list","of","float","`","Tensor","`","current","values","of","parameters","being","optimized",".","auxiliary_state",":","the","(","structure","of",")","`","Tensor","`","(","s",")","containing","state","carried","forward","from","the","previous","step",".","Returns",":","has_converged",":","boolean","`","Tensor","`","indicating","whether","the","optimization","has","converged",".","updated_auxiliary_state",":","(","Structure","of",")","`","Tensor","`","(","s",")","representing","updated","quantities","tracked","by","the","convergence","criterion",".","This","should","match","the","structure","of","the","value","returned","by","`","bootstrap","`","."]
code: def one_step(self, step, loss, grads, parameters, auxiliary_state):
      """Updates tracked quantities for a new step, and determines if converged.
  
      The shape of the `Tensor`s specifying `loss`, `grads`, and `parameters` may
      optionally be prefixed by one or more batch dimension(s). In this case,
      the returned value `has_converged` will have shape equal to the broadcast
      batch shape of whichever of those quantities is used by this convergence
      criterion, and the quantities defining the convergence criterion (
      `min_num_steps`, etc.).
  
      Args:
        step: integer `Tensor` index of the current step, where `step >= 1` (on
          step `0`, `initial_state` should be called instead).
        loss: float `Tensor`  value of loss at the current step.
        grads: list of float `Tensor` gradients of `loss` wrt `parameters`.
        parameters: list of float `Tensor` current values of parameters
          being optimized.
        auxiliary_state: the (structure of) `Tensor`(s) containing state carried
          forward from the previous step.
      Returns:
        has_converged: boolean `Tensor` indicating whether the optimization has
          converged.
        updated_auxiliary_state: (Structure of) `Tensor`(s) representing
          updated quantities tracked by the convergence criterion. This should
          match the structure of the value returned by `bootstrap`.
      """
      with tf.name_scope(self.name):
          has_converged, updated_auxiliary_state = self._one_step(
              step=tf.convert_to_tensor(step),
              loss=tf.convert_to_tensor(loss),
              grads=(
                  tf.nest.map_structure(tf.convert_to_tensor, grads)
                  if grads is not None
                  else grads
              ),
              parameters=tf.nest.map_structure(
                  tf.convert_to_tensor, parameters
              ),
              auxiliary_state=auxiliary_state,
          )
          if self.min_num_steps is not None:
              has_converged = has_converged & (
                  step >= self.min_num_steps
              )
          return has_converged, updated_auxiliary_state
  
file_path: probability/tensorflow_probability/python/optimizer/convergence_criteria/convergence_criterion.py
imports: import tensorflow.compat.v2
label_desc: Use image node_id 5 for calling the ConvergenceCriterion obj's underlying member method code with example usage: obj.one_step(step, loss, grads, parameters, auxiliary_state) and returns: has_converged,  updated_auxiliary_state