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: Checks that input partition dims are valid for the `Tensor`.
  
  Args:
    tensor: Input tensor for partitioning.
    dims: A list of integer describes how to partition the input tensor.
  
  Raises:
    ValueError: If the tensor can't be partitioned by dims or the
      num_cores_per_replica doesn't match the number of
      partitions(dims.prod()).
code: def _check_input_partition_dims(self, tensor, dims):
      """Checks that input partition dims are valid for the `Tensor`.
  
      Args:
        tensor: Input tensor for partitioning.
        dims: A list of integer describes how to partition the input tensor.
  
      Raises:
        ValueError: If the tensor can't be partitioned by dims or the
          num_cores_per_replica doesn't match the number of
          partitions(dims.prod()).
      """
      # No partitioning specified, so don't perform further checks.
      if dims is None:
          return
  
      dims = np.array(dims)
  
      if (dims < 1).any():
          raise ValueError("All input partition dims must be >= 1.")
  
      # No partitioning, so don't perform further checks.
      if dims.prod() == 1:
          return
  
      if dims.prod() != self._device_assignment.num_cores_per_replica:
          raise ValueError(
              "The product of each input partition dim should equal to "
              "num_cores_per_replica. (dim = {}, num_cores_per_replica "
              "= {})".format(
                  dims, self._device_assignment.num_cores_per_replica
              )
          )
      if dims.shape[0] != tensor.shape.ndims:
          raise ValueError(
              "Input partition dims must have the same number of dimensions "
              "as the `Tensor` to be partitioned. (tensor shape = {}, input "
              "partition dims = {}).".format(
                  tensor.shape.as_list(), dims
              )
          )
  
      tensor.shape.assert_is_fully_defined()
  
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 4 for calling the _PartitionedInfeedQueue obj's underlying member method code with example usage: obj._check_input_partition_dims(tensor, dims) without return types