// RUN: tf-opt %s -tfl-legalize-tf='preserve-assert-op=true' | FileCheck %s --dump-input=fail

func.func @preserve_assert(%arg0: tensor<1xi32>, %arg1: tensor<1xi32>) -> tensor<1xi1> {
  %0 = "tf.LessEqual"(%arg0, %arg1) : (tensor<1xi32>, tensor<1xi32>) -> tensor<1xi1>
  "tf.Assert"(%0, %arg1) {summarize = 3} : (tensor<1xi1>, tensor<1xi32>) -> ()
  func.return %0 : tensor<1xi1>
  // CHECK-LABEL: preserve_assert
  // CHECK: tfl.less_equal
  // CHECK: Assert

// RUN: odml-to-stablehlo-opt %s -stablehlo-custom-call-legalize-composite | FileCheck %s

// CHECK-LABEL: module
module {
  // CHECK-LABEL: @main
  func.func @main(%arg0: tensor<1xf32>, %arg1: tensor<2xf32>) {
    // CHECK: stablehlo.custom_call @foo
    stablehlo.custom_call @foo() : () -> ()
    // CHECK-NOT: stablehlo.custom_call

// RUN: tf-opt %s -tfl-prepare-tf -tfl-legalize-tf='run-tfl-runtime-verification=false'  -tfl-optimize | FileCheck %s

func.func @broadcast_to_bf16(%arg0: tensor<3xbf16>, %arg1: tensor<2xi64>) -> tensor<3x3xbf16> {
  %0 = "tf.BroadcastTo"(%arg0, %arg1) : (tensor<3xbf16>, tensor<2xi64>) -> tensor<3x3xbf16>
  func.return %0: tensor<3x3xbf16>

// CHECK-LABEL: broadcast_to_bf16
// CHECK:  [[CST:%.*]] = arith.constant dense<1.000000e+00> : tensor<3x3xbf16>

// RUN: tf-opt -xla-legalize-tf %s | FileCheck %s

//===----------------------------------------------------------------------===//
// tf.BatchMatMulV2 op legalizations.
//===----------------------------------------------------------------------===//

func.func @batchmatmulv2_basic(%arg0: tensor<1x4x2xf32>, %arg1: tensor<3x2x4xf32>) -> tensor<3x4x4xf32> {
// CHECK-LABEL:   func @batchmatmulv2_basic

// RUN: tf-opt -tfl-legalize-tensorlist -canonicalize -split-input-file %s | FileCheck %s

// CHECK-LABEL: listReserveScalarShapeI32
func.func @listReserveScalarShapeI32(%arg0: tensor<i32>, %arg1: tensor<i32>) -> tensor<!tf_type.variant<tensor<*xi32>>> {
  %0 = "tf.TensorListReserve"(%arg0, %arg1) : (tensor<i32>, tensor<i32>) -> tensor<!tf_type.variant<tensor<*xi32>>>

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

// RUN: tf-opt %s -split-input-file -tfl-legalize-hashtables-tf --cse | FileCheck %s

// Test for case with string -> int64 hashtable.
func.func @hashtable_string_to_int64(%arg0: tensor<i64>) -> tensor<*xi64> {
  %cst = arith.constant dense<"f"> : tensor<!tf_type.string>

        continue;
      if (auto size_op = llvm::dyn_cast<TF::LookupTableSizeV2Op>(user))
        continue;

      return false;
    }
  }
  return true;
}

// Pass which legalizes TF hash tables only when they are covered by the
// TensorFlow Lite hash table kernels.
class LegalizeHashTablesPass
    : public impl::LegalizeHashTablesPassBase<LegalizeHashTablesPass> {
 public:

                                         llvm::SmallVector<Value>& outputs);

  // Tries to legalize the specified TensorFlow op, if supported.
  //
  // Emits an error and returns failure if an error is encountered during
  // conversion. Note that success return value doesn't mean successful
  // legalization.
  mlir::LogicalResult LegalizeOp();

  // Converts the given operand to expression of kind kConstant or kXlaOp.

  ];
}

def LegalizeHloToTfPass : Pass<"tf-legalize-hlo", "ModuleOp"> {
  let summary = "Legalize from MHLO to the TF dialect";
  let dependentDialects = ["TF::TensorFlowDialect"];
  let constructor = "mlir::odml::CreateLegalizeHloToTfPass()";
}

def LegalizeHloToTfLitePass : Pass<"tfl-legalize-hlo", "mlir::ModuleOp"> {
  let summary = "Legalize from MHLO to the TFLite dialect";