int rank = tensor_type.getRank();
  if (rank <= partition_dim) {
    return op->emitError() << "cannot partition " << first_operand_type
                           << " (rank = " << rank << ") along dimension "
                           << partition_dim << ".";
  }

  llvm::SmallVector<int64_t, 4> partition_dims(is_replicated ? 0 : rank, 1);
  if (!is_replicated) {
    partition_dims[partition_dim] = num_cores_per_replica.getInt();
  }

  // CHECK-SAME: _XlaSharding = ""
  // CHECK-SAME: partition_dims = []
  // CHECK: [[PI_1:%.*]] = "tf.TPUPartitionedInputV2"([[ARG2]], [[ARG3]])
  // CHECK-SAME: _XlaSharding = ""
  // CHECK-SAME: partition_dims = []
  // CHECK: [[RI:%.*]] = "tf.TPUReplicatedInput"([[PI_0]], [[PI_1]])

// CHECK-NEXT: = "tf.Const"
// CHECK-NEXT: = "tf.AddV2"
    %outputs, %control = tf_executor.island wraps "tf.TPUPartitionedInput"(%arg0, %arg1) {_XlaSharding = "\08\03\1A\02\02\01\22\02\00\01", device = "", partition_dim = 0 : i64} : (tensor<2x4xf32>, tensor<2x4xf32>) -> tensor<4x4xf32>

      %pi_0 = "tf.TPUPartitionedInput"(%arg0, %arg1) {_XlaSharding = "", device = "", partition_dim = -1 : i64} : (!rtype, !rtype) -> !rtype
      %pi_1 = "tf.TPUPartitionedInput"(%arg2, %arg3) {_XlaSharding = "", device = "", partition_dim = -1 : i64} : (!rtype, !rtype) -> !rtype
      %ri = "tf.TPUReplicatedInput"(%pi_0, %pi_1) : (!rtype, !rtype) -> !rtype
      return %ri : !rtype

  %pi_0 = "tf.TPUPartitionedInput"(%arg0, %arg1) {_XlaSharding = "", device = "", partition_dim = -1 : i64} : (!rtype, !rtype) -> !rtype
  %pi_1 = "tf.TPUPartitionedInput"(%arg2, %arg3) {_XlaSharding = "", device = "", partition_dim = -1 : i64} : (!rtype, !rtype) -> !rtype
  %ri = "tf.TPUReplicatedInput"(%pi_0, %pi_1) : (!rtype, !rtype) -> !rtype
  return %ri : !rtype
}
```

  let summary = [{
An op that groups a list of partitioned inputs together. This op
  }];

  let arguments = (ins
    Variadic<TF_Tensor>:$inputs,

    DefaultValuedOptionalAttr<I64Attr, "0">:$partition_dim,
    OptionalAttr<StrAttr>:$_XlaSharding
  );

  let results = (outs
    TF_Tensor:$output
  );

  TF_DerivedOperandTypeAttr T = TF_DerivedOperandTypeAttr<0>;

  // CHECK:      [[INPUT:%.+]] = "tf.TPUPartitionedInputV2"([[READ0]], [[READ1]])
  // CHECK-SAME: _XlaSharding = ""
  // CHECK-SAME: partition_dims = []
  %0 = "tf.TPUPartitionedInputV2"(%arg0, %arg1) {N = 2 : i64, _XlaSharding = "", partition_dims = []} : (tensor<!tf_type.resource<tensor<i32>>>, tensor<!tf_type.resource<tensor<i32>>>) -> tensor<!tf_type.resource<tensor<i32>>>

  %pi_0 = "tf.TPUPartitionedInputV2"(%arg0, %arg1) {_XlaSharding = "", partition_dims = []} : (tensor<!tf_type.resource<tensor<10x3xf32>>>, tensor<!tf_type.resource<tensor<10x3xf32>>>) -> tensor<!tf_type.resource<tensor<10x3xf32>>>
  // expected-error@+1 {{expects partition_dims = [] but found [1, 2]}}

    // Abort if TPUPartitionedInputV2(s) do not have the same attributes.
    if (!llvm::equal(partition_dims, op_partition_dims)) {
      return partitioned_input->emitOpError()
             << "expects partition_dims = " << partition_dims << " but found "
             << op_partition_dims;
    } else if (partitioned_input.getIsPacked() !=
               first_partitioned_input.getIsPacked()) {

  %1:2 = "tf_device.cluster_func"(%0, %arg1) {func = @cluster_func, use_spmd_for_xla_partitioning = true, num_cores_per_replica = 1 : i64} : (tensor<*xi32>, tensor<*xi32>) -> (tensor<*xi32>, tensor<*xi32>)
  %2 = "tf.TPUPartitionedOutputV2"(%1#1) {_XlaSharding = "\08\01\1A\01\01\22\01\00", partition_dims = []} : (tensor<*xi32>) -> tensor<*xi32>