[JAX SC] Update MergeAll method to take an absl::Span<PartitionedCooTensors>. This allows for more flexible usage and avoids unnecessary moves.

PiperOrigin-RevId: 845459258

Author: Google-ML-Automation

jax_tpu_embedding/sparsecore/lib/core/partitioned_coo_tensors.h

@@ -207,7 +207,7 @@ class PartitionedCooTensors {
   int GetNumMinibatches() const { return bucket_count_per_sc_; }
 
   static PartitionedCooTensors MergeAll(
-      std::vector<PartitionedCooTensors>&& parts) {
+      absl::Span<PartitionedCooTensors> parts) {
     DCHECK(!parts.empty());
     // If there is only one part, no merging is needed.
     if (parts.size() == 1) {

jax_tpu_embedding/sparsecore/lib/core/partitioned_coo_tensors_test.cc

@@ -19,6 +19,7 @@
 
 #include <gmock/gmock.h>
 #include <gtest/gtest.h>
+#include "absl/types/span.h"  // from @com_google_absl
 #include "jax_tpu_embedding/sparsecore/lib/core/coo_format.h"
 
 namespace jax_sc_embedding {
@@ -217,7 +218,7 @@ TEST(PartitionedCooTensorsTest, MergeDistinctSparseCores) {
   parts.push_back(std::move(part2));
 
   PartitionedCooTensors result =
-      PartitionedCooTensors::MergeAll(std::move(parts));
+      PartitionedCooTensors::MergeAll(absl::MakeSpan(parts));
 
   // Verify that the result contains 2 SCs.
   EXPECT_EQ(result.GetNumMinibatches(), 2);
@@ -259,7 +260,7 @@ TEST(PartitionedCooTensorsTest, MergeDistinctSparseCoresWithOffsets) {
   parts.push_back(std::move(part2));
 
   PartitionedCooTensors result =
-      PartitionedCooTensors::MergeAll(std::move(parts));
+      PartitionedCooTensors::MergeAll(absl::MakeSpan(parts));
 
   EXPECT_EQ(result.GetNumMinibatches(), 2);
   // SC 0 has 2 buckets, sizes 1 and 1.

jax_tpu_embedding/sparsecore/lib/core/sort_and_group_coo_tensors_impl.h

@@ -523,7 +523,7 @@ SortAndGroupCooTensorsPerLocalDeviceImpl(
       // Merge the parts for this task.
       SparseCoreTaskResult task_result = {
           .grouped_coo_tensors =
-              PartitionedCooTensors::MergeAll(std::move(sc_parts)),
+              PartitionedCooTensors::MergeAll(absl::MakeSpan(sc_parts)),
           .stats_host = std::move(local_stats_host),
           .dropped_id_count = total_dropped,
           .split_val = task_split};
@@ -575,7 +575,7 @@ SortAndGroupCooTensorsPerLocalDeviceImpl(
         }
 
         PartitionedCooTensors merged =
-            PartitionedCooTensors::MergeAll(std::move(parts));
+            PartitionedCooTensors::MergeAll(absl::MakeSpan(parts));
         final_result.emplace(std::make_pair(std::move(merged), total_dropped));
       });