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: Get internal state for the given TensorFlow device.
docstr_tok: ["Get","internal","state","for","the","given","TensorFlow","device","."]
code: def _get_device(self, device_name: str):
      """Get internal state for the given TensorFlow device."""
      tfutil.assert_tf_initialized()
      if device_name in self._devices:
          return self._devices[device_name]
  
      # Initialize fields.
      device = util.EasyDict()
      device.name = device_name
      device.optimizer = (
          None  # Underlying optimizer:     optimizer_class
      )
      device.loss_scaling_var = (
          None  # Log2 of loss scaling:     tf.Variable
      )
      device.grad_raw = (
          OrderedDict()
      )  # Raw gradients:            var => [grad, ...]
      device.grad_clean = (
          OrderedDict()
      )  # Clean gradients:          var => grad
      device.grad_acc_vars = (
          OrderedDict()
      )  # Accumulation sums:        var => tf.Variable
      device.grad_acc_count = (
          None  # Accumulation counter:     tf.Variable
      )
      device.grad_acc = (
          OrderedDict()
      )  # Accumulated gradients:    var => grad
  
      # Setup TensorFlow objects.
      with tfutil.absolute_name_scope(
          self.scope + "/Devices"
      ), tf.device(device_name), tf.control_dependencies(None):
          if device_name not in self._shared_optimizers:
              optimizer_name = self.scope.replace(
                  "/", "_"
              ) + "_opt%d" % len(self._shared_optimizers)
              self._shared_optimizers[
                  device_name
              ] = self.optimizer_class(
                  name=optimizer_name,
                  learning_rate=self.learning_rate,
                  **self.optimizer_kwargs,
              )
          device.optimizer = self._shared_optimizers[device_name]
          if self.use_loss_scaling:
              device.loss_scaling_var = tf.Variable(
                  np.float32(self.loss_scaling_init),
                  trainable=False,
                  name="loss_scaling_var",
              )
  
      # Register device.
      self._devices[device_name] = device
      return device
  
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 2 for calling the Optimizer obj's underlying member method code with example usage: obj._get_device(device_name) and returns: device, self