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: Register the gradients of the given loss function with respect to the given variables.
  Intended to be called once per GPU.
docstr_tok: ["Register","the","gradients","of","the","given","loss","function","with","respect","to","the","given","variables",".","Intended","to","be","called","once","per","GPU","."]
code: def register_gradients(
      self, loss: TfExpression, trainable_vars: Union[List, dict]
  ) -> None:
      """Register the gradients of the given loss function with respect to the given variables.
      Intended to be called once per GPU."""
      tfutil.assert_tf_initialized()
      assert not self._updates_applied
      device = self._get_device(loss.device)
  
      # Validate trainables.
      if isinstance(trainable_vars, dict):
          trainable_vars = list(
              trainable_vars.values()
          )  # allow passing in Network.trainables as vars
      assert (
          isinstance(trainable_vars, list) and len(trainable_vars) >= 1
      )
      assert all(
          tfutil.is_tf_expression(expr)
          for expr in trainable_vars + [loss]
      )
      assert all(var.device == device.name for var in trainable_vars)
  
      # Validate shapes.
      if self._gradient_shapes is None:
          self._gradient_shapes = [
              var.shape.as_list() for var in trainable_vars
          ]
      assert len(trainable_vars) == len(self._gradient_shapes)
      assert all(
          var.shape.as_list() == var_shape
          for var, var_shape in zip(
              trainable_vars, self._gradient_shapes
          )
      )
  
      # Report memory usage if requested.
      deps = []
      if self._report_mem_usage:
          self._report_mem_usage = False
          try:
              with tf.name_scope(self.id + "_mem"), tf.device(
                  device.name
              ), tf.control_dependencies([loss]):
                  deps.append(
                      autosummary.autosummary(
                          self.id + "/mem_usage_gb",
                          tf.contrib.memory_stats.BytesInUse()
                          / 2**30,
                      )
                  )
          except tf.errors.NotFoundError:
              pass
  
      # Compute gradients.
      with tf.name_scope(self.id + "_grad"), tf.device(
          device.name
      ), tf.control_dependencies(deps):
          loss = self.apply_loss_scaling(tf.cast(loss, tf.float32))
          gate = (
              tf.train.Optimizer.GATE_NONE
          )  # disable gating to reduce memory usage
          grad_list = device.optimizer.compute_gradients(
              loss=loss, var_list=trainable_vars, gate_gradients=gate
          )
  
      # Register gradients.
      for grad, var in grad_list:
          if var not in device.grad_raw:
              device.grad_raw[var] = []
          device.grad_raw[var].append(grad)
  
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 3 for calling the Optimizer obj's underlying member method code with example usage: obj.register_gradients(loss, trainable_vars) without return types