}

    auto partitioned_output = builder.create<TF::TPUPartitionedOutputV2Op>(
        cluster_func->getLoc(), partitioned_output_types, result,
        partitioned_input.getPartitionDimsAttr(),
        partitioned_input.get_XlaShardingAttr());
    for (auto [i, value] : llvm::enumerate(partitioned_output.getOutput())) {
      const auto& resource = packed_input ? inputs[0] : inputs[i];

      }
    }

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