If pre-populated reduced groups are present, make sure to copy over t…

…he reduced intervals as well. PiperOrigin-RevId: 630408147
tensorflow · May 3, 2024 · 44bd31b · 44bd31b
1 parent 90b100e
commit 44bd31b
Showing 1 changed file with 5 additions and 3 deletions.
diff --git a/third_party/xla/xla/hlo/experimental/auto_sharding/auto_sharding_solver.cc b/third_party/xla/xla/hlo/experimental/auto_sharding/auto_sharding_solver.cc
@@ -257,10 +257,12 @@ std::optional<std::pair<int64_t, int64_t>> ReduceMemoryTerms(
             : reducer.Reduce(num_lives, num_primitives, std::move(Intervals));
     reduced_intervals = reducer.GetReducedIntervals();
     reduced_groups = reducer.GetReducedGroups();
-  } else {
+  } else {  // If we've already done term reduction, just copy over the results.
+    for (const auto& interval : intervals) {
+      reduced_intervals.push_back({interval.first(), interval.second()});
+    }
     for (const auto& group : groups) {
-      reduced_groups.push_back(
-          absl::btree_set<int64_t>(group.prims().begin(), group.prims().end()));
+      reduced_groups.push_back({group.prims().begin(), group.prims().end()});
     }
   }
   solver.MakeIntVarArray(reduced_groups.size(), 0.0, MPSolver::infinity(),