[XLA] Make HloModule::computations() return raw pointers.

Like HloComputation::instructions(), HloModule::computations() used to
return a list of unique_ptrs.  But this is an implementation detail that
shouldn't be leaked into the public API.

This patch also adds HloModule::MakeNonFusionComputations(), because
many of the callers of computations() went on to filter out all the
fusion computations.

It would be possible to implement MakeNonFusionComputations() "in place"
using a filtering iterator, but I don't think it's necessary -- we never
have *that* many computations, and since many callers go on to copy the
list of non-fusion computations, making it unconditionally a copy is
simpler and avoids a footgun.

PiperOrigin-RevId: 170529051

1 files changed, 1 insertions, 8 deletions

diff --git a/tensorflow/compiler/xla/service/transpose_folding.cc b/tensorflow/compiler/xla/service/transpose_folding.cc
index d668c812f4..816c8a7485 100644
--- a/tensorflow/compiler/xla/service/transpose_folding.cc
+++ b/tensorflow/compiler/xla/service/transpose_folding.cc
@@ -171,14 +171,7 @@ StatusOr<bool> TransposeFolding::Run(HloModule* module) {
     return tensorflow::Status::OK();
   };
 
-  std::vector<HloComputation*> computations;
-  for (auto& computation : module->computations()) {
-    if (computation->IsFusionComputation()) {
-      continue;
-    }
-    computations.push_back(computation.get());
-  }
-  for (auto& comp : computations) {
+  for (auto* comp : module->MakeNonfusionComputations()) {
     TF_RETURN_IF_ERROR(comp->Accept(visit_fn));
   }