class_docstr: A Wrapper for tf.train.Optimizer.
  
  Automatically takes care of:
  - Gradient averaging for multi-GPU training.
  - Gradient accumulation for arbitrarily large minibatches.
  - Dynamic loss scaling and typecasts for FP16 training.
  - Ignoring corrupted gradients that contain NaNs/Infs.
  - Reporting statistics.
  - Well-chosen default settings.
class_docstr_tok: ["A","Wrapper","for","tf.train.Optimizer",".","Automatically","takes","care","of",":","-","Gradient","averaging","for","multi-GPU","training",".","-","Gradient","accumulation","for","arbitrarily","large","minibatches",".","-","Dynamic","loss","scaling","and","typecasts","for","FP16","training",".","-","Ignoring","corrupted","gradients","that","contain","NaNs\/Infs",".","-","Reporting","statistics",".","-","Well-chosen","default","settings","."]
docstr: Construct training op to update the registered variables based on their gradients.
docstr_tok: ["Construct","training","op","to","update","the","registered","variables","based","on","their","gradients","."]
code: def apply_updates(self, allow_no_op: bool = False) -> tf.Operation:
      """Construct training op to update the registered variables based on their gradients."""
      tfutil.assert_tf_initialized()
      assert not self._updates_applied
      self._updates_applied = True
      all_ops = []
  
      # Check for no-op.
      if allow_no_op and len(self._devices) == 0:
          with tfutil.absolute_name_scope(self.scope):
              return tf.no_op(name="TrainingOp")
  
      # Clean up gradients.
      for device_idx, device in enumerate(self._devices.values()):
          with tfutil.absolute_name_scope(
              self.scope + "/Clean%d" % device_idx
          ), tf.device(device.name):
              for var, grad in device.grad_raw.items():
                  # Filter out disconnected gradients and convert to float32.
                  grad = [g for g in grad if g is not None]
                  grad = [tf.cast(g, tf.float32) for g in grad]
  
                  # Sum within the device.
                  if len(grad) == 0:
                      grad = tf.zeros(
                          var.shape
                      )  # No gradients => zero.
                  elif len(grad) == 1:
                      grad = grad[0]  # Single gradient => use as is.
                  else:
                      grad = tf.add_n(
                          grad
                      )  # Multiple gradients => sum.
  
                  # Scale as needed.
                  scale = (
                      1.0
                      / len(device.grad_raw[var])
                      / len(self._devices)
                  )
                  scale = tf.constant(
                      scale, dtype=tf.float32, name="scale"
                  )
                  if self.minibatch_multiplier is not None:
                      scale /= tf.cast(
                          self.minibatch_multiplier, tf.float32
                      )
                  scale = self.undo_loss_scaling(scale)
                  device.grad_clean[var] = grad * scale
  
      # Sum gradients across devices.
      if len(self._devices) > 1:
          with tfutil.absolute_name_scope(
              self.scope + "/Broadcast"
          ), tf.device(None):
              for all_vars in zip(
                  *[
                      device.grad_clean.keys()
                      for device in self._devices.values()
                  ]
              ):
                  if len(all_vars) > 0 and all(
                      dim > 0 for dim in all_vars[0].shape.as_list()
                  ):  # NCCL does not support zero-sized tensors.
                      all_grads = [
                          device.grad_clean[var]
                          for device, var in zip(
                              self._devices.values(), all_vars
                          )
                      ]
                      all_grads = nccl_ops.all_sum(all_grads)
                      for device, var, grad in zip(
                          self._devices.values(), all_vars, all_grads
                      ):
                          device.grad_clean[var] = grad
  
      # Apply updates separately on each device.
      for device_idx, device in enumerate(self._devices.values()):
          with tfutil.absolute_name_scope(
              self.scope + "/Apply%d" % device_idx
          ), tf.device(device.name):
              # pylint: disable=cell-var-from-loop
  
              # Accumulate gradients over time.
              if self.minibatch_multiplier is None:
                  acc_ok = tf.constant(True, name="acc_ok")
                  device.grad_acc = OrderedDict(device.grad_clean)
              else:
                  # Create variables.
                  with tf.control_dependencies(None):
                      for var in device.grad_clean.keys():
                          device.grad_acc_vars[var] = tf.Variable(
                              tf.zeros(var.shape),
                              trainable=False,
                              name="grad_acc_var",
                          )
                      device.grad_acc_count = tf.Variable(
                          tf.zeros([]),
                          trainable=False,
                          name="grad_acc_count",
                      )
  
                  # Track counter.
                  count_cur = device.grad_acc_count + 1.0
                  count_inc_op = lambda: tf.assign(
                      device.grad_acc_count, count_cur
                  )
                  count_reset_op = lambda: tf.assign(
                      device.grad_acc_count, tf.zeros([])
                  )
                  acc_ok = count_cur >= tf.cast(
                      self.minibatch_multiplier, tf.float32
                  )
                  all_ops.append(
                      tf.cond(acc_ok, count_reset_op, count_inc_op)
                  )
  
                  # Track gradients.
                  for var, grad in device.grad_clean.items():
                      acc_var = device.grad_acc_vars[var]
                      acc_cur = acc_var + grad
                      device.grad_acc[var] = acc_cur
                      with tf.control_dependencies([acc_cur]):
                          acc_inc_op = lambda: tf.assign(
                              acc_var, acc_cur
                          )
                          acc_reset_op = lambda: tf.assign(
                              acc_var, tf.zeros(var.shape)
                          )
                          all_ops.append(
                              tf.cond(acc_ok, acc_reset_op, acc_inc_op)
                          )
  
              # No overflow => apply gradients.
              all_ok = tf.reduce_all(
                  tf.stack(
                      [acc_ok]
                      + [
                          tf.reduce_all(tf.is_finite(g))
                          for g in device.grad_acc.values()
                      ]
                  )
              )
              apply_op = lambda: device.optimizer.apply_gradients(
                  [
                      (tf.cast(grad, var.dtype), var)
                      for var, grad in device.grad_acc.items()
                  ]
              )
              all_ops.append(tf.cond(all_ok, apply_op, tf.no_op))
  
              # Adjust loss scaling.
              if self.use_loss_scaling:
                  ls_inc_op = lambda: tf.assign_add(
                      device.loss_scaling_var, self.loss_scaling_inc
                  )
                  ls_dec_op = lambda: tf.assign_sub(
                      device.loss_scaling_var, self.loss_scaling_dec
                  )
                  ls_update_op = lambda: tf.group(
                      tf.cond(all_ok, ls_inc_op, ls_dec_op)
                  )
                  all_ops.append(
                      tf.cond(acc_ok, ls_update_op, tf.no_op)
                  )
  
              # Last device => report statistics.
              if device_idx == len(self._devices) - 1:
                  all_ops.append(
                      autosummary.autosummary(
                          self.id + "/learning_rate", self.learning_rate
                      )
                  )
                  all_ops.append(
                      autosummary.autosummary(
                          self.id + "/overflow_frequency",
                          tf.where(all_ok, 0, 1),
                          condition=acc_ok,
                      )
                  )
                  if self.use_loss_scaling:
                      all_ops.append(
                          autosummary.autosummary(
                              self.id + "/loss_scaling_log2",
                              device.loss_scaling_var,
                          )
                      )
  
      # Initialize variables.
      self.reset_optimizer_state()
      if self.use_loss_scaling:
          tfutil.init_uninitialized_vars(
              [
                  device.loss_scaling_var
                  for device in self._devices.values()
              ]
          )
      if self.minibatch_multiplier is not None:
          tfutil.run(
              [
                  var.initializer
                  for device in self._devices.values()
                  for var in list(device.grad_acc_vars.values())
                  + [device.grad_acc_count]
              ]
          )
  
      # Group everything into a single op.
      with tfutil.absolute_name_scope(self.scope):
          return tf.group(*all_ops, name="TrainingOp")
  
imports: import numpy
  import tensorflow
  from collections import OrderedDict
  from typing import List, Union
  from .None import autosummary
  from .None import tfutil
  from ..None import util
  from .tfutil import TfExpression, TfExpressionEx
label_desc: Use image node_id 4 for calling the Optimizer obj's underlying member method code with example usage: obj.apply_updates(allow_no_op) and returns: tf, tf