class_docstr: A helper object to build a device infeed queue with input partition.
  
  Args:
    number_of_tuple_elements: the number of Tensors fed atomically through the
      queue, must be present unless it can be inferred from other arguments.
    device_assignment: A TPU `DeviceAssignment` which is used to place all the
      partitions to different TPU infeed queues.
    host_id: The id of the host machine.
    input_partition_dims: A nested list/tuple of integers. Each inner
      list/tuple describes how to partition the corresponding input tensor.
    tuple_types: If not None, a list of types of the elements of the queue.
    tuple_shapes: If not None, a list of shapes of the elements of the queue.
    name: The name of the queue.
class_docstr_tok: ["A","helper","object","to","build","a","device","infeed","queue","with","input","partition",".","Args",":","number_of_tuple_elements",":","the","number","of","Tensors","fed","atomically","through","the","queue",",","must","be","present","unless","it","can","be","inferred","from","other","arguments",".","device_assignment",":","A","TPU","`","DeviceAssignment","`","which","is","used","to","place","all","the","partitions","to","different","TPU","infeed","queues",".","host_id",":","The","id","of","the","host","machine",".","input_partition_dims",":","A","nested","list\/tuple","of","integers",".","Each","inner","list\/tuple","describes","how","to","partition","the","corresponding","input","tensor",".","tuple_types",":","If","not","None",",","a","list","of","types","of","the","elements","of","the","queue",".","tuple_shapes",":","If","not","None",",","a","list","of","shapes","of","the","elements","of","the","queue",".","name",":","The","name","of","the","queue","."]
docstr: Generates the host-side Ops to enqueue the partitioned inputs.
  
  sharded_inputs is a list, one for each replica, of lists of
  Tensors. sharded_inputs[i] is the tuple of Tensors to use to feed
  replica i.
  sharded_inputs[i][j] is partitioned by self._input_partition_dims[j].
  
  For example, if sharded_inputs[i][j] is a 2-D Tensor:
  [[A, B, C, D],
   [E ,F, G, H]]
  self._input_partition_dims[j] is [2, 4].
  
  sharded_inputs[i][j] will be partitioned and flattened into:
  [A, B, C, D, E, F, G, H] and fed into the logical core ids:
  [0, 1, 2, 3, 4, 5, 6, 7] respectively.
  
  Args:
    sharded_inputs: a list of lists of Tensors. The length of the
      outer list determines the number of shards. Each inner list indicates
      the types and shapes of the tuples in the corresponding shard.
  
  Returns:
    A list of host-side Ops, one for each shard, that when executed together
    will enqueue a full-size element of infeed.
  
  Raises:
    ValueError: if the queue configuration has previously been frozen and the
      shapes of the elements of sharded_inputs are not compatible with the
      frozen configuration; or if the shapes of the elements of sharded_inputs
      don't form a consistent unsharded tuple; or if the elements of a tuple
      have different device constraints; or if the partition dims are invalid.
    TypeError: if the queue configuration has previously been frozen and the
      types of the elements of sharded_inputs are not compatible with the
      frozen configuration; or if the types of the elements of sharded_inputs
      don't form a consistent unsharded tuple.
code: def generate_enqueue_ops(
      self, sharded_inputs
  ):  # pytype: disable=signature-mismatch  # overriding-parameter-count-checks
      """Generates the host-side Ops to enqueue the partitioned inputs.
  
      sharded_inputs is a list, one for each replica, of lists of
      Tensors. sharded_inputs[i] is the tuple of Tensors to use to feed
      replica i.
      sharded_inputs[i][j] is partitioned by self._input_partition_dims[j].
  
      For example, if sharded_inputs[i][j] is a 2-D Tensor:
      [[A, B, C, D],
       [E ,F, G, H]]
      self._input_partition_dims[j] is [2, 4].
  
      sharded_inputs[i][j] will be partitioned and flattened into:
      [A, B, C, D, E, F, G, H] and fed into the logical core ids:
      [0, 1, 2, 3, 4, 5, 6, 7] respectively.
  
      Args:
        sharded_inputs: a list of lists of Tensors. The length of the
          outer list determines the number of shards. Each inner list indicates
          the types and shapes of the tuples in the corresponding shard.
  
      Returns:
        A list of host-side Ops, one for each shard, that when executed together
        will enqueue a full-size element of infeed.
  
      Raises:
        ValueError: if the queue configuration has previously been frozen and the
          shapes of the elements of sharded_inputs are not compatible with the
          frozen configuration; or if the shapes of the elements of sharded_inputs
          don't form a consistent unsharded tuple; or if the elements of a tuple
          have different device constraints; or if the partition dims are invalid.
        TypeError: if the queue configuration has previously been frozen and the
          types of the elements of sharded_inputs are not compatible with the
          frozen configuration; or if the types of the elements of sharded_inputs
          don't form a consistent unsharded tuple.
      """
      self.set_configuration_from_sharded_input_tensors(sharded_inputs)
      number_of_replicas = len(sharded_inputs)
      number_of_tuple_elements = len(sharded_inputs[0])
  
      assert len(self._input_partition_dims) == number_of_tuple_elements
      enqueue_ops = []
  
      for replica_index in range(number_of_replicas):
          flattened_inputs = sharded_inputs[replica_index]
          inputs_part_dims_flat = nest.flatten_up_to(
              flattened_inputs, self._input_partition_dims
          )
          inputs_parted_iters = [
              iter(
                  self._check_dims_and_partition_or_replicate_on_host(
                      x, dims
                  )
              )
              for x, dims in zip(
                  sharded_inputs[replica_index], inputs_part_dims_flat
              )
          ]
  
          # Find the replica_id of the host's logical core 0.
          # The self._host_id is guaranteed to contain the logical core 0,
          # even when num_cores_per_replica > num_cores_per_host -- the function
          # caller makes sure that this host_id will must be receiving data (calls
          # input_fn).
          replica_id = self._device_assignment.lookup_replicas(
              task_id=self._host_id, logical_core=0
          )[replica_index]
          for logical_core in range(
              self._device_assignment.num_cores_per_replica
          ):
              # Places different partitions to different logic cores.
              # Since there can be multiple hosts per replica, we need to find
              # the actual host (device) of this logical core.
              device = self._device_assignment.host_device(
                  replica=replica_id, logical_core=logical_core
              )
  
              with ops.device(device):
                  ordinal = self._device_assignment.tpu_ordinal(
                      replica=replica_id, logical_core=logical_core
                  )
                  infeed_inputs = []
                  for it in inputs_parted_iters:
                      input_for_device = next(it, None)
                      if input_for_device is not None:
                          infeed_inputs.append(input_for_device)
  
                  if infeed_inputs:
                      enqueue_ops.append(
                          tpu_ops.infeed_enqueue_tuple(
                              inputs=infeed_inputs,
                              shapes=[x.shape for x in infeed_inputs],
                              name="enqueue/replica_{0}/input_{1}".format(
                                  replica_index, logical_core
                              ),
                              device_ordinal=ordinal,
                          )
                      )
      return enqueue_ops
  
imports: import itertools
  import numpy
  from tensorflow.python.compiler.xla.experimental import xla_sharding
  from tensorflow.python.framework import dtypes
  from tensorflow.python.framework import ops
  from tensorflow.python.framework import tensor_shape
  from tensorflow.python.ops import array_ops
  from tensorflow.python.tpu import tpu_name_util
  from tensorflow.python.tpu import tpu_sharding
  from tensorflow.python.tpu.ops import tpu_ops
  from tensorflow.python.util import nest
label_desc: Use image node_id 3 for calling the _PartitionedInfeedQueue obj's underlying member method code with example usage: obj.generate_enqueue_ops(sharded_inputs) and returns: enqueue_ops