}
    }

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