Replace the use of xla::OkStatus with absl::OkStatus now that they're…

… the same.

PiperOrigin-RevId: 635902683

third_party/xla/xla/hlo/experimental/auto_sharding/auto_sharding.cc

@@ -3289,7 +3289,7 @@ absl::Status GenerateReduceScatter(
             .output_sharding);
     TF_RETURN_IF_ERROR(inst->ReplaceAllUsesWith(replace_with));
   }
-  return OkStatus();
+  return absl::OkStatus();
 }
 
 void AnnotateShardingWithSimpleHeuristic(
@@ -3447,7 +3447,7 @@ absl::Status FilterStrategy(const HloInstruction* ins, const Shape& shape,
       << ins->ToString() << " does not have any valid strategies";
   strategy_group->strategies = std::move(new_strategies);
 
-  return OkStatus();
+  return absl::OkStatus();
 }
 
 // Return the output sharding of the reduce-scatter variant of a given strategy.
@@ -3744,7 +3744,7 @@ absl::Status AutoShardingImplementation::CanonicalizeLayouts(
     HloModule* module) {
   if (!module->layout_canonicalization_callback()) {
     LOG(INFO) << "There is no registered layout_canonicalization_callback.";
-    return OkStatus();
+    return absl::OkStatus();
   }
   TF_ASSIGN_OR_RETURN(auto layouts,
                       module->layout_canonicalization_callback()(*module));
@@ -3763,7 +3763,7 @@ absl::Status AutoShardingImplementation::CanonicalizeLayouts(
   }
   *module->mutable_config().mutable_entry_computation_layout() =
       entry_computation_layout;
-  return OkStatus();
+  return absl::OkStatus();
 }
 
 // Computes the set of instructions that lie outside any manually partitioned

third_party/xla/xla/hlo/experimental/auto_sharding/auto_sharding_dot_handler.cc

@@ -26,6 +26,7 @@ limitations under the License.
 #include "absl/container/flat_hash_set.h"
 #include "absl/log/check.h"
 #include "absl/log/log.h"
+#include "absl/status/status.h"
 #include "absl/strings/match.h"
 #include "absl/strings/str_cat.h"
 #include "absl/strings/str_format.h"
@@ -1035,7 +1036,7 @@ absl::Status DotHandler::RegisterStrategies() {
                                       cluster_env_, batch_map_, option_));
   }
 
-  return OkStatus();
+  return absl::OkStatus();
 }
 
 /************** ConvHandler function definitions **************/
@@ -1106,7 +1107,7 @@ absl::Status ConvHandler::RegisterStrategies() {
                                       cluster_env_, batch_map_, option_));
   }
 
-  return OkStatus();
+  return absl::OkStatus();
 }
 
 void ConvHandler::SplitLhsBatchRhsOutchannel() {
@@ -1233,7 +1234,7 @@ absl::Status HandleDot(std::unique_ptr<StrategyGroup>& strategy_group,
                      instruction_id, instruction_sequence, hlo_cost_analysis,
                      cluster_env, batch_map, option, call_graph);
   TF_RETURN_IF_ERROR(handler.RegisterStrategies());
-  return OkStatus();
+  return absl::OkStatus();
 }
 
 // Register strategies for convolution instructions.
@@ -1266,7 +1267,7 @@ absl::Status HandleConv(std::unique_ptr<StrategyGroup>& strategy_group,
     TF_RETURN_IF_ERROR(handler.RegisterStrategies());
   }
 
-  return OkStatus();
+  return absl::OkStatus();
 }
 
 }  // namespace spmd

third_party/xla/xla/hlo/experimental/auto_sharding/auto_sharding_runner.cc

@@ -17,6 +17,7 @@ limitations under the License.
 #include <ostream>
 #include <string>
 
+#include "absl/status/status.h"
 #include "xla/hlo/experimental/auto_sharding/auto_sharding.h"
 #include "xla/hlo/ir/hlo_module.h"
 #include "xla/service/hlo_parser.h"
@@ -44,7 +45,7 @@ absl::Status RunAutoShardingPassFromFile(const std::string& file_name) {
   TF_ASSIGN_OR_RETURN(bool changed, AutoSharding(option).Run(hlo_module.get()));
   CHECK(changed);
   std::cout << hlo_module->ToString() << std::endl;
-  return OkStatus();
+  return absl::OkStatus();
 }
 
 }  // namespace

third_party/xla/xla/hlo/experimental/auto_sharding/auto_sharding_solver.cc

@@ -38,6 +38,7 @@ limitations under the License.
 #include "absl/container/flat_hash_set.h"
 #include "absl/log/check.h"
 #include "absl/log/log.h"
+#include "absl/status/status.h"
 #include "absl/strings/str_cat.h"
 #include "absl/time/clock.h"
 #include "absl/time/time.h"
@@ -1268,7 +1269,7 @@ absl::Status ValidateRequest(const AutoShardingSolverRequest& request) {
     const int num_v_strategies = request.computation_costs(v).costs_size();
     CHECK_EQ(num_strategies, num_u_strategies * num_v_strategies);
   }
-  return OkStatus();
+  return absl::OkStatus();
 }
 
 }  // namespace spmd

third_party/xla/xla/hlo/experimental/auto_sharding/auto_sharding_util.cc

@@ -1859,7 +1859,7 @@ absl::Status CheckAliasSetCompatibility(const AliasSet& alias_set,
       }
     }
   }
-  return OkStatus();
+  return absl::OkStatus();
 }
 
 size_t VectorGreaterThanOneElementCount(absl::Span<const int64_t> span,