tf.raw_ops.TPUPartitionedInputV2

An op that groups a list of partitioned inputs together. Supports ND sharding.

inputs A list of at least 1 Tensor objects with the same type. A list of partitioned inputs which must have the same shape.
partition_dims A list of ints. A list of integers describing how each dimension is partitioned. Emptiness indicates the inputs are replicated.
is_packed An optional bool. Defaults to False. Indicates whether the input is a packed resource.
name A name for the operation (optional).

A Tensor. Has the same type as inputs.