"@local_tsl//tsl/platform:errors",
        "@local_xla//xla:shape_util",
        "@local_xla//xla:xla_data_proto_cc",
        "@local_xla//xla/client:xla_computation",
        "@local_xla//xla/hlo/ir:hlo",
        "@local_xla//xla/mlir_hlo",
        "@local_xla//xla/mlir_hlo:hlo_dialect_registration",
        "@local_xla//xla/mlir_hlo:mhlo_passes",
        "@local_xla//xla/translate/mhlo_to_hlo:layout_util",

    std::vector<std::unique_ptr<mlir::Pass>>& custom_legalization_passes,
    xla::CompileOnlyClient* client, XlaCompilationResult* compilation_result) {
  LOG_FIRST_N(INFO, 1) << "Compiling MLIR computation to XLA HLO using the "
                          "Combined MLIR Tf2Xla Bridge.";

  absl::StatusOr<std::string> mlir_compilation =
      internal::CompileFromMlirToXlaHlo(

#include "tensorflow/compiler/jit/xla_compilation_cache.pb.h"
#include "tensorflow/compiler/jit/xla_device_compiler_client.h"
#include "tensorflow/compiler/tf2xla/xla_compiler.h"
#include "xla/pjrt/pjrt_client.h"
#include "xla/service/hlo.pb.h"
#include "xla/util.h"
#include "tensorflow/core/framework/device.h"
#include "tensorflow/core/lib/strings/proto_serialization.h"
#include "tensorflow/core/platform/path.h"
#include "tensorflow/core/platform/status.h"

// appropriate executable isn't found on disk, compiles the given Tensorflow
// function/graph/cluster into an XlaCompilationResult (HLO) and
// `ExecutableType` and tries saving/persisting the compiled HLO and executable
// to disk.
//
// Since XLA computations must have static shapes, DeviceCompiler generates a
// new XLA computation for each new set of input shapes.

  // inside mlir is done.
  if ((import_saved_model_object_graph || import_saved_model_signature_defs) &&
      import_hlo) {
    llvm::errs() << "Import saved model and import hlo cannot be both set.";
    return kTrFailure;
  }

  if (import_saved_model_object_graph || import_saved_model_signature_defs) {
    // Saved model import path.
    int saved_model_version;

// The full pipeline of converting jax random include 2 steps.
// 1. Rename the jax random functions to tflite wrapped functions with the aid
//    of "jax.named_call". For example, in the dumped hlo, the
//    jax.random.uniform will have name "tfl_wrapped_jax_random_uniform".
// 2. Replace the body of "tfl_wrapped_jax_random_uniform" and
//    "tfl_wrapped_jax_random_normal" with tfl.CustomOp("RandomUniform") and

        full set of deps needed by the generated library.
      enable_xla_hlo_profiling: Enable XLA HLO profiling in the generated
        program, and emit metadata that lets us pretty-print the gathered
        profile counters.
      enable_tracemes: Tell tfcompile to generate calls to
        TraceMe::Activity{Start|End} around HLO instructions that can be used by
        Xprof to construct profiler timelines.

// RUN: odml-to-stablehlo-opt %s -stablehlo-composite-legalize-tfl-custom | FileCheck %s
// RUN: tf_tfl_translate --enable-hlo-to-tf-conversion --input-mlir %s -o - | flatbuffer_translate --tflite-flatbuffer-to-mlir - -o - | FileCheck %s --check-prefix=CHECK-ROUNDTRIP

module {
  func.func public @main(%arg0: tensor<3x3xf32>, %arg1: tensor<3x3xf32>, %arg2: tensor<1x100x32x4xf32>,

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

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

#include <numeric>
#include <string>
#include <utility>

#include "llvm/ADT/StringRef.h"
#include "mlir/Dialect/Func/IR/FuncOps.h"  // from @llvm-project

#include "tensorflow/compiler/jit/tests/device_compiler_test_helper.h"

#include <string>

#include "absl/strings/match.h"
#include "tensorflow/compiler/jit/xla_compilation_cache.pb.h"
#include "xla/service/hlo.pb.h"
#include "tensorflow/core/platform/path.h"
#include "tensorflow/core/public/session.h"

namespace tensorflow {
namespace {

// Creates a float tensor of linearly increasing values, starting from offset.