builder.getI64ArrayAttr(partition_dims), builder.getBoolAttr(false),
        op.get_XlaShardingAttr());
    op->replaceAllUsesWith(pi);
  } else {
    auto po = builder.create<TF::TPUPartitionedOutputV2Op>(
        op.getLoc(), op.getResultTypes(), op.getOperand(),
        builder.getI64ArrayAttr(partition_dims), op.get_XlaShardingAttr());
    op->replaceAllUsesWith(po);
  }

  return success();
}

                             llvm::SmallVector<Value, 4> new_reads,
                             bool is_packed) {
  xla::OpSharding sharding;
  sharding.ParseFromString(
      old_partitioned_input.get_XlaShardingAttr().getValue().str());
  for (OpOperand& read_use :
       llvm::make_early_inc_range(old_read.getValue().getUses())) {
    if (dyn_cast_or_null<tf_device::ClusterFuncOp>(read_use.getOwner())) {

    // Assume all the replicas have the same structure.
    mlir::TF::TPUPartitionedOutputV2Op first_op = *(ops.begin());
    mlir::ArrayAttr dims = first_op.getPartitionDimsAttr();
    StringAttr sharding = first_op.get_XlaShardingAttr();
    Operation::result_type_range output_types = first_op.getResultTypes();
    result_op = builder.create<mlir::TF::TPUPartitionedOutputV2Op>(

        data_and_token->setAttr(
            kShardingAttr,
            rewriter.getStringAttr(sharding_proto.SerializeAsString()));
      } else {
        data_and_token->setAttr(kShardingAttr, op.get_XlaShardingAttr());
      }
    }

    if (op->hasAttr("layouts")) {
      // Append a UnitAttr for the "token" operand of the mhlo.infeed op here to