// tensor to sparse format.
std::unique_ptr<OperationPass<func::FuncOp>> CreateDenseToSparsePass();

// Creates function pass to legalize TF While to TFL While.
std::unique_ptr<OperationPass<ModuleOp>> CreateLegalizeTFWhilePass();

// Legalize tflite flex ops to TF ops.
std::unique_ptr<OperationPass<func::FuncOp>> CreateLiftTfliteFlexOpsPass();

  ];
}

def LegalizeWhilePass : Pass<"tfl-legalize-tf-while", "mlir::ModuleOp"> {
  let summary = "Legalize from TensorFlow While to TensorFlow Lite While.";
  let constructor = "CreateLegalizeTFWhilePass()";
  let dependentDialects = ["TFL::TensorFlowLiteDialect"];
}

def LegalizeVariablesPass : Pass<"tfl-legalize-variables-tf", "mlir::ModuleOp"> {
  let summary = "Legalize TensorFlow variables to TensorFlow Lite dialect.";

// RUN: tf-opt --tfl-legalize-tf-while %s -o - | flatbuffer_translate -mlir-to-tflite-flatbuffer - -o -  | flatbuffer_translate --tflite-flatbuffer-to-mlir - -o - | FileCheck %s

// CHECK-LABEL: main
func.func @main(%arg0: tensor<*xf32>) -> tensor<*xf32> {
  // CHECK:     %{{.*}} = tfl.add %{{.*}}, %{{.*}} {fused_activation_function = "NONE"} : tensor<*xf32>
  // CHECK:     return %{{.*}} : tensor<*xf32>

        "prepare-tf.mlir": "medium",
        "prepare-tf-fake-quant.mlir": "medium",
        "prepare-tf-fake-quant-4bit.mlir": "medium",
        "raise-custom-ops.mlir": "medium",
    },
    tags_override = {
        "legalize-tf.mlir": ["no_rocm"],
        "optimize.mlir": ["no_rocm"],
        "prepare-tf.mlir": ["no_rocm"],
        "const-fold.mlir": ["no_mac_arm64"],
    },
    test_file_exts = ["mlir"],
)

    pass_manager->addPass(mlir::TF::CreateTFShapeInferencePass());
  }

  // Legalize while early to allow further constant folding.
  // TODO(jpienaar): This may not actually matter as we do canonicalization
  // after the legalize below, for now it needs to be below the above passes
  // that work on TF dialect and before inliner so that the function calls in

def VerifyQuantLegalization : Pass<"verify-quant-legalization", "mlir::func::FuncOp"> {
  let summary = "Verifies that all TF quant ops and types have been legalized.";

  let description = [{
    Ensures that all TF quant ops and types have been legalized to HLO
    and reports an error about which op failed to legalize. This pass
    does not transform any ops and is checking.}];

  let constructor = "mlir::quant::stablehlo::CreateVerifyQuantLegalizationPass()";

    ],
    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>> {

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

}

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*/ {});