pass_manager.addPass(
      mlir::odml::CreateLegalizeTFXlaCallModuleToStablehloPass());

  // Legalize MHLO to StableHLO should be moved closer to where it is needed
  // There are some entry points that start with HLO->MHLO like
  // jax_to_tfl_flatbuffer.cc which can likely be updated to emit StableHLO
  // to be consistent with other entrypoints.
  pass_manager.addPass(mlir::mhlo::createHloLegalizeToStablehloPass());

#include "tensorflow/compiler/aot/tests/test_graph_tfvariable_sequential_updates_mlir_bridge.h"
// Similarly, there are files for testing the MLIR based lowering of HLO to
// object code for XLA:CPU
#elif defined(MHLO_LOWERING_TEST)
#include "tensorflow/compiler/aot/tests/test_graph_tfadd_mhlo_lowering.h"
#include "tensorflow/compiler/aot/tests/test_graph_tfadd_with_ckpt_mhlo_lowering.h"

        "@local_tsl//tsl/platform:statusor",
        "@local_xla//xla:xla_data_proto_cc",
        "@local_xla//xla/client:xla_builder",
        "@local_xla//xla/client:xla_computation",
        "@local_xla//xla/hlo/ir:hlo",
        "@local_xla//xla/mlir_hlo",
        "@local_xla//xla/service:hlo_proto_cc",
        "@local_xla//xla/translate/hlo_to_mhlo:hlo_module_importer",
        "@local_xla//xla/translate/hlo_to_mhlo:hlo_to_mlir_hlo",

// This is the legalization pattern definition file for HLO to TF.

include "mlir/IR/OpBase.td"
include "mlir/Dialect/Func/IR/FuncOps.td"
include "stablehlo/dialect/ChloOps.td"
include "tensorflow/compiler/mlir/tensorflow/ir/tf_ops.td"
include "mhlo/IR/hlo_ops.td"

// Check if broadcasting is compatible with TF ops.
def IsLegalNumpyRankedBroadcast :

#include "tensorflow/compiler/tf2xla/xla_compiler.h"
#include "tensorflow/compiler/tf2xla/xla_helpers.h"
#include "xla/client/local_client.h"
#include "xla/executable_run_options.h"
#include "xla/hlo/ir/hlo_input_output_alias_config.h"
#include "xla/pjrt/pjrt_client.h"
#include "xla/pjrt/tf_pjrt_client.h"
#include "xla/service/executable.h"
#include "xla/service/gpu/gpu_executable_run_options.h"

  // have `xla_device_metadata`. This is used for TFRT-TPU when
  // BuildXlaDeviceCompiler() is called in GetCompilerIr(). Currently only
  // lowering to HLO is needed there and xla::LocalClient doesn't support
  // building the executable for TFRT-TPU and hence, is set to nullptr here.
  if (platform_info.device_type() == DEVICE_TPU) {
    *xla_device_compiler = CreateXlaDeviceCompiler(

  // Convert model.
  if (enable_mlir_converter) {
    if (model_flags.use_hlo_import() && model_flags.has_saved_model_dir()) {
      PyErr_SetString(PyExc_ValueError,
                      "Cannot specify both saved_model and hlo import.");
      return nullptr;
    }

    if (model_flags.use_hlo_import()) {
      status = tensorflow::ConvertJaxToTFLiteFlatBuffer(
          input_contents_txt, model_flags, toco_flags,

        "@llvm-project//mlir:Support",
        "@local_xla//xla:status_macros",
        "@local_xla//xla:xla_data_proto_cc",
        "@local_xla//xla/client:sharding_builder",
        "@local_xla//xla/hlo/ir:hlo",
        "@local_xla//xla/service:hlo_parser",
    ],
)

cc_library(
    name = "location_utils",
    srcs = ["utils/location_utils.cc"],
    hdrs = ["utils/location_utils.h"],
    deps = [

limitations under the License.
==============================================================================*/

// This file implements logic for lowering TensorFlow dialect's communication
// ops (TF/XLA) to the HLO dialect.

#include <atomic>
#include <cstddef>
#include <cstdint>
#include <memory>
#include <optional>
#include <string>

#include "llvm/ADT/DenseMap.h"
#include "llvm/ADT/STLExtras.h"

      Env::Default(), GetFilePath(key2, persistor.persistent_cache_directory()),
      entry));

  MockXlaCompilerClient mock_client;
  // Try to load executable corresponding to key2 (whose file contains HLO
  // corresponding to key1).
  auto loaded_executable = persistor.TryToLoadExecutable(
      /*signature_hash=*/123, "signature", DefaultXlaOptions(),
      compilation_result_mul, &mock_client);