class_docstr: Optimizer that runs in parallel across multiple local devices.
  
  TFMultiGPUTowerStack automatically splits up and loads training data
  onto specified local devices (e.g. GPUs) with `load_data()`. During a call
  to `optimize()`, the devices compute gradients over slices of the data in
  parallel. The gradients are then averaged and applied to the shared
  weights.
  
  The data loaded is pinned in device memory until the next call to
  `load_data`, so you can make multiple passes (possibly in randomized order)
  over the same data once loaded.
  
  This is similar to tf1.train.SyncReplicasOptimizer, but works within a
  single TensorFlow graph, i.e. implements in-graph replicated training:
  
  https://www.tensorflow.org/api_docs/python/tf/train/SyncReplicasOptimizer
docstr: Bulk loads the specified inputs into device memory.
  
  The shape of the inputs must conform to the shapes of the input
  placeholders this optimizer was constructed with.
  
  The data is split equally across all the devices. If the data is not
  evenly divisible by the batch size, excess data will be discarded.
  
  Args:
      sess: TensorFlow session.
      inputs: List of arrays matching the input placeholders, of shape
          [BATCH_SIZE, ...].
      state_inputs: List of RNN input arrays. These arrays have size
          [BATCH_SIZE / MAX_SEQ_LEN, ...].
      num_grad_updates: The lifetime number of gradient updates that the
          policy having collected the data has already undergone.
  
  Returns:
      The number of tuples loaded per device.
docstr_tok: ["Bulk","loads","the","specified","inputs","into","device","memory",".","The","shape","of","the","inputs","must","conform","to","the","shapes","of","the","input","placeholders","this","optimizer","was","constructed","with",".","The","data","is","split","equally","across","all","the","devices",".","If","the","data","is","not","evenly","divisible","by","the","batch","size",",","excess","data","will","be","discarded",".","Args",":","sess",":","TensorFlow","session",".","inputs",":","List","of","arrays","matching","the","input","placeholders",",","of","shape","[","BATCH_SIZE",",","...","]",".","state_inputs",":","List","of","RNN","input","arrays",".","These","arrays","have","size","[","BATCH_SIZE","\/","MAX_SEQ_LEN",",","...","]",".","num_grad_updates",":","The","lifetime","number","of","gradient","updates","that","the","policy","having","collected","the","data","has","already","undergone",".","Returns",":","The","number","of","tuples","loaded","per","device","."]
code: def load_data(
      self, sess, inputs, state_inputs, num_grad_updates=None
  ):
      """Bulk loads the specified inputs into device memory.
  
      The shape of the inputs must conform to the shapes of the input
      placeholders this optimizer was constructed with.
  
      The data is split equally across all the devices. If the data is not
      evenly divisible by the batch size, excess data will be discarded.
  
      Args:
          sess: TensorFlow session.
          inputs: List of arrays matching the input placeholders, of shape
              [BATCH_SIZE, ...].
          state_inputs: List of RNN input arrays. These arrays have size
              [BATCH_SIZE / MAX_SEQ_LEN, ...].
          num_grad_updates: The lifetime number of gradient updates that the
              policy having collected the data has already undergone.
  
      Returns:
          The number of tuples loaded per device.
      """
      self.num_grad_updates = num_grad_updates
  
      if log_once("load_data"):
          logger.info(
              "Training on concatenated sample batches:\n\n{}\n".format(
                  summarize(
                      {
                          "placeholders": self.loss_inputs,
                          "inputs": inputs,
                          "state_inputs": state_inputs,
                      }
                  )
              )
          )
  
      feed_dict = {}
      assert len(self.loss_inputs) == len(inputs + state_inputs), (
          self.loss_inputs,
          inputs,
          state_inputs,
      )
  
      # Let's suppose we have the following input data, and 2 devices:
      # 1 2 3 4 5 6 7                              <- state inputs shape
      # A A A B B B C C C D D D E E E F F F G G G  <- inputs shape
      # The data is truncated and split across devices as follows:
      # |---| seq len = 3
      # |---------------------------------| seq batch size = 6 seqs
      # |----------------| per device batch size = 9 tuples
  
      if len(state_inputs) > 0:
          smallest_array = state_inputs[0]
          seq_len = len(inputs[0]) // len(state_inputs[0])
          self._loaded_max_seq_len = seq_len
      else:
          smallest_array = inputs[0]
          self._loaded_max_seq_len = 1
  
      sequences_per_minibatch = (
          self.max_per_device_batch_size
          // self._loaded_max_seq_len
          * len(self.devices)
      )
      if sequences_per_minibatch < 1:
          logger.warning(
              (
                  "Target minibatch size is {}, however the rollout sequence "
                  "length is {}, hence the minibatch size will be raised to "
                  "{}."
              ).format(
                  self.max_per_device_batch_size,
                  self._loaded_max_seq_len,
                  self._loaded_max_seq_len * len(self.devices),
              )
          )
          sequences_per_minibatch = 1
  
      if len(smallest_array) < sequences_per_minibatch:
          # Dynamically shrink the batch size if insufficient data
          sequences_per_minibatch = _make_divisible_by(
              len(smallest_array), len(self.devices)
          )
  
      if log_once("data_slicing"):
          logger.info(
              (
                  "Divided {} rollout sequences, each of length {}, among "
                  "{} devices."
              ).format(
                  len(smallest_array),
                  self._loaded_max_seq_len,
                  len(self.devices),
              )
          )
  
      if sequences_per_minibatch < len(self.devices):
          raise ValueError(
              "Must load at least 1 tuple sequence per device. Try "
              "increasing `sgd_minibatch_size` or reducing `max_seq_len` "
              "to ensure that at least one sequence fits per device."
          )
      self._loaded_per_device_batch_size = (
          sequences_per_minibatch
          // len(self.devices)
          * self._loaded_max_seq_len
      )
  
      if len(state_inputs) > 0:
          # First truncate the RNN state arrays to the sequences_per_minib.
          state_inputs = [
              _make_divisible_by(arr, sequences_per_minibatch)
              for arr in state_inputs
          ]
          # Then truncate the data inputs to match
          inputs = [
              arr[: len(state_inputs[0]) * seq_len] for arr in inputs
          ]
          assert len(state_inputs[0]) * seq_len == len(inputs[0]), (
              len(state_inputs[0]),
              sequences_per_minibatch,
              seq_len,
              len(inputs[0]),
          )
          for ph, arr in zip(self.loss_inputs, inputs + state_inputs):
              feed_dict[ph] = arr
          truncated_len = len(inputs[0])
      else:
          truncated_len = 0
          for ph, arr in zip(self.loss_inputs, inputs):
              truncated_arr = _make_divisible_by(
                  arr, sequences_per_minibatch
              )
              feed_dict[ph] = truncated_arr
              if truncated_len == 0:
                  truncated_len = len(truncated_arr)
  
      sess.run([t.init_op for t in self._towers], feed_dict=feed_dict)
  
      self.num_tuples_loaded = truncated_len
      samples_per_device = truncated_len // len(self.devices)
      assert samples_per_device > 0, "No data loaded?"
      assert (
          samples_per_device % self._loaded_per_device_batch_size == 0
      )
      # Return loaded samples per-device.
      return samples_per_device
  
imports: from collections import namedtuple, OrderedDict
  import gymnasium
  import logging
  import re
  import tree
  from typing import Callable, Dict, List, Optional, Tuple, Type, Union
  from ray.util.debug import log_once
  from ray.rllib.models.tf.tf_action_dist import TFActionDistribution
  from ray.rllib.models.modelv2 import ModelV2
  from ray.rllib.policy.policy import Policy
  from ray.rllib.policy.sample_batch import SampleBatch
  from ray.rllib.policy.tf_policy import TFPolicy
  from ray.rllib.policy.view_requirement import ViewRequirement
  from ray.rllib.models.catalog import ModelCatalog
  from ray.rllib.utils import force_list
  from ray.rllib.utils.annotations import override, DeveloperAPI
  from ray.rllib.utils.debug import summarize
  from ray.rllib.utils.deprecation import deprecation_warning, DEPRECATED_VALUE
  from ray.rllib.utils.framework import try_import_tf
  from ray.rllib.utils.metrics import DIFF_NUM_GRAD_UPDATES_VS_SAMPLER_POLICY, NUM_GRAD_UPDATES_LIFETIME
  from ray.rllib.utils.spaces.space_utils import get_dummy_batch_for_space
  from ray.rllib.utils.tf_utils import get_placeholder
  from ray.rllib.utils.typing import LocalOptimizer, ModelGradients, TensorType, AlgorithmConfigDict
label_desc: Use image node_id 2 for calling the TFMultiGPUTowerStack obj's underlying member method code with example usage: obj.load_data(sess, inputs, state_inputs, num_grad_updates) and returns: samples_per_device