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
code: def _setup_device(
      self, tower_i, device, device_input_placeholders, num_data_in
  ):
      assert num_data_in <= len(device_input_placeholders)
      with tf.device(device):
          with tf1.name_scope(TOWER_SCOPE_NAME + f"_{tower_i}"):
              device_input_batches = []
              device_input_slices = []
              for i, ph in enumerate(device_input_placeholders):
                  current_batch = tf1.Variable(
                      ph,
                      trainable=False,
                      validate_shape=False,
                      collections=[],
                  )
                  device_input_batches.append(current_batch)
                  if i < num_data_in:
                      scale = self._max_seq_len
                      granularity = self._max_seq_len
                  else:
                      scale = self._max_seq_len
                      granularity = 1
                  current_slice = tf.slice(
                      current_batch,
                      (
                          [self._batch_index // scale * granularity]
                          + [0] * len(ph.shape[1:])
                      ),
                      (
                          [
                              self._per_device_batch_size
                              // scale
                              * granularity
                          ]
                          + [-1] * len(ph.shape[1:])
                      ),
                  )
                  current_slice.set_shape(ph.shape)
                  device_input_slices.append(current_slice)
              graph_obj = self.policy_copy(device_input_slices)
              device_grads = graph_obj.gradients(
                  self.optimizers, graph_obj._losses
              )
          return _Tower(
              tf.group(
                  *[batch.initializer for batch in device_input_batches]
              ),
              device_grads,
              graph_obj,
          )
  
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 5 for calling the TFMultiGPUTowerStack obj's underlying member method code with example usage: obj._setup_device(tower_i, device, device_input_placeholders, num_data_in) and returns: _Tower