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();
  }

      %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>>>

  %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>

func.func @partitioned_outputs(%arg0: tensor<?xi32>) -> (tensor<?xi32>, tensor<?xi32>, tensor<?xi32>, tensor<?xi32>) {
  %pi0 = "tf.TPUPartitionedInputV2"(%arg0) {N = 2, partition_dims = [], _XlaSharding = "", is_packed = true} : (tensor<?xi32>) -> (tensor<?xi32>)
  %pi1 = "tf.TPUPartitionedInputV2"(%arg0) {N = 2, partition_dims = [], _XlaSharding = "", is_packed = true} : (tensor<?xi32>) -> (tensor<?xi32>)

    %read1 = "tf.ReadVariableOp"(%arg1) : (tensor<!tf_type.resource<tensor<i32>>>) -> tensor<i32>
    // CHECK-NOT: tf.TPUPartitionedInputV2
    %partitioned_input = "tf.TPUPartitionedInputV2"(%read0, %read1) {N = 2 : i64, partition_dims = []} : (tensor<i32>, tensor<i32>) -> tensor<i32>
    // CHECK: %[[COMPILE_OUTPUT:[0-9]*]]:3 = "tf_device.launch"
    // CHECK-NEXT: "tf._TPUCompileMlir"()
    // CHECK: "tf_device.launch"

  TPUPartitionedInputV2Op op = *this;

  int num_partitions = 1;
  const mlir::ArrayAttr partition_dims = op.getPartitionDims();
  for (const mlir::Attribute &dim : partition_dims) {
    num_partitions *= dim.cast<IntegerAttr>().getInt();
  }

  const bool is_packed = op.getIsPacked();
  const bool replicated = partition_dims.empty();
  const int num_inputs_expected = is_packed ? 1 : num_partitions;

^bb0(%arg0: tensor<2x4xf32>, %arg1: tensor<2x4xf32>):
  // expected-error @+1 {{expected 1 inputs, got 2}}
  %0 = "tf.TPUPartitionedInputV2"(%arg0, %arg1) {partition_dims = [2, 1], is_packed = true} : (tensor<2x4xf32>, tensor<2x4xf32>) -> tensor<4x4xf32>
  func.return %0 : tensor<4x4xf32>
}

// -----

// Test invalid tf.TPUPartitionedInputV2 without packing