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]])

    // 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()) {