],
    licenses = ["notice"],
)

glob_lit_tests(
    name = "all_tests",
    data = [":test_utilities"],
    driver = "//tensorflow/compiler/mlir/lite/stablehlo:run_lit.sh",
    size_override = {
        "legalize-skip-quantization-ops.mlir": "medium",
    },
    test_file_exts = [
        "mlir",
        "cc",
    ],
)

# Bundle together all of the test utilities that are used by tests.
filegroup(

namespace mlir {
namespace TF {

//===----------------------------------------------------------------------===//
// The pass to legalize the quantization emulation ops from TF.
//
namespace {

// Legalize TF quantization emulation ops to that in Quant ops dialect.
struct LegalizeTFToQuant
    : public PassWrapper<LegalizeTFToQuant, OperationPass<func::FuncOp>> {

- The Canonicalization pass iteratively applies canonicalization
transformations in a greedy way until no further changes occur.
Canonicalization includes constant folding.
- The Legalize pass converts TensorFlow operations to TensorFlow Lite
ones. The operations that cannot be mapped to TensorFlow Lite dialect
are left as TensorFlow operations. Unsupported op handling follows the
proposed TFLite mechanism.

// RUN: tf-opt -tfl-legalize-random %s | FileCheck %s


// CHECK-LABEL:   func @tfl_wrapped_jax_random_normal(
// CHECK-SAME:                                        %[[RNG:.*]]: tensor<2xui32>) -> tuple<tensor<3x4xf32>> {
// CHECK:           %[[VAL_0:.*]] = stablehlo.constant dense<[3, 4]> : tensor<2xi32>

// RUN: tf-opt -tfe-legalize-tfg %s | FileCheck %s

module  {
  tfg.graph #tf_type.version<producer = 27, min_consumer = 0> {
    // CHECK: tf_executor.Enter
    // CHECK: {{%.*}}, %[[TOKEN:.*]], {{%.*}} = tf_executor.NextIteration.Source
    // CHECK: {{%.*}}, {{%.*}}, %[[CONTROL:.*]] = tf_executor.Merge
    // CHECK: tf_executor.island(%[[CONTROL]]) wraps "tf.Const"()
    // CHECK: tf_executor.LoopCond
    // CHECK: tf_executor.Switch

}

TEST(CompileMlirUtil, HasLegalizationPass) {
  OpPassManager pass_manager;
  llvm::StringRef device_type = "XLA_CPU_JIT";
  absl::string_view kLegalizeTfPass = "xla-legalize-tf";

  CreateConvertMlirToXlaHloPipeline(pass_manager, device_type,
                                    /*enable_op_fallback=*/true,
                                    /*custom_legalization_passes*/ {});

auto* mlir_failed_legalization_op_count =
    tensorflow::monitoring::Counter<1>::New(
        "/tensorflow/core/tf2xla/"
        "mlir_second_phase_failed_legalization_op_count",
        "Counts which op fails to legalize", "op_name");

auto* mlir_non_static_op_count = tensorflow::monitoring::Counter<1>::New(
    "/tensorflow/core/tf2xla/"
    "mlir_second_phase_non_static_op_count",

See the License for the specific language governing permissions and
limitations under the License.
==============================================================================*/

// Legalize mhlo.dot_general to tflite.batch_matmul.

#ifndef TENSORFLOW_COMPILER_MLIR_LITE_STABLEHLO_TRANSFORMS_LEGALIZE_HLO_CONVERSIONS_DOT_GENERAL_H_

#include "xla/client/compile_only_client.h"
#include "tensorflow/core/tpu/kernels/tpu_compile_op_support.h"
#include "tsl/platform/statusor.h"

namespace tensorflow {
namespace tf2xla {
namespace internal {

// Legalize the given MLIR module to XLA HLO using a combination of the MLIR
// Bridge and XlaBuilder
absl::StatusOr<XlaCompilationResult> LegalizeTfToHlo(
    const tpu::MlirToHloArgs& computation,

  }

  if (failed(mlir::tf_saved_model::FreezeVariables(
          module_op, saved_model_bundle->GetSession()))) {
    return absl::InternalError("Failed to freeze variables.");
  }

  // Run legalize TF to StableHLO pass to convert `tf.Const` and
  // `tf.Const`->`tf.Cast` patterns after variable freezing. The TF shape
  // inference pass is also required to resolve unknown shapes in the TF dialect