};

// Converts tf.RandomShuffle op into a series of XLA HLO ops.
//
// tf.RandomShuffle shuffles tensors along the first dimension. If the input
// tensor's rank is 1, then it is translated into HLO sort op(s) according to
// indices randomly generated via HLO rng_uniform ops. Otherwise, it is
// translated into an HLO while op to first emulate shuffling indices using
// HLO dynamic_slice and dynamic_update_slice ops, then finally HLO gather

config.name = 'MLIR ' + os.path.basename(config.mlir_test_dir)

config.test_format = lit.formats.ShTest(not llvm_config.use_lit_shell)

# suffixes: A list of file extensions to treat as test files.
config.suffixes = ['.cc', '.hlo', '.json', '.mlir', '.pbtxt', '.py']

# test_source_root: The root path where tests are located.
config.test_source_root = config.mlir_test_dir

# test_exec_root: The root path where tests should be run.

#include "tensorflow/compiler/tf2xla/layout_util.h"
#include "tensorflow/compiler/tf2xla/xla_compiler.h"
#include "tensorflow/compiler/tf2xla/xla_helpers.h"
#include "xla/client/compile_only_client.h"
#include "xla/hlo/ir/hlo_input_output_alias_config.h"
#include "xla/mlir_hlo/mhlo/IR/register.h"
#include "xla/shape.h"
#include "xla/shape_util.h"
#include "xla/status_macros.h"
#include "tensorflow/core/framework/function.h"

        "@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",

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

tf_cc_test(
    name = "tf_mlir_translate_registration_test",
    size = "small",
    srcs = ["tf_mlir_translate_registration_test.cc"],

        "@local_tsl//tsl/platform:status",
        "@local_tsl//tsl/platform:statusor",
        "@local_xla//xla:xla_proto_cc",
        "@local_xla//xla/client:compile_only_client",
        "@local_xla//xla/hlo/ir:hlo",
        "@local_xla//xla/mlir_hlo:hlo_dialect_registration",
        "@local_xla//xla/pjrt:compile_options_proto_cc",
        "@stablehlo//:register",
    ],
)

tf_cc_test(
    name = "legalize_tf_test",

  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());

  // 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,

  Type new_arg_type = tensorflow::GetTypeFromTFTensorShape(shape, element_type);

  if (auto input_ty = mlir::dyn_cast<RankedTensorType>(old_arg_type)) {
    ArrayRef<int64_t> bounds = hlo::encodingToBounds(input_ty.getEncoding());
    // The input type has bounded dynamic dimension.
    if (!bounds.empty()) {
      SmallVector<int64_t> new_bounds(bounds.begin(), bounds.end());

  let description = [{
    This pass looks for the usage of the result of TPUCopyWithDynamicShapeOp
    and sets the shape of these inputs to be dynamic shaped. This will ensure
    that the generated HLO program is correctly reflecting the dynamic shape.
  }];

  // Required for mhlo bounded shape extension.
  let dependentDialects = ["mhlo::MhloDialect"];