func.func @TfAdd(%arg0: tensor<4xf64>, %arg1: tensor<4xf64>) -> tensor<4xf64> {
  %0 = "tfl.custom"(%arg0, %arg1) {
    custom_code = "FlexAdd",
    custom_option = #tfl<const_bytes : "0x03416464001412034164641A001A002A070A015412023002320000021B171414042801">
  } : (tensor<4xf64>, tensor<4xf64>) -> tensor<4xf64>

// CHECK: "tf.Add"(%arg0, %arg1) {T = f64}  : (tensor<4xf64>, tensor<4xf64>) -> tensor<4xf64>
  func.return %0 : tensor<4xf64>
}

// RUN: flatbuffer_translate -mlir-to-tflite-flatbuffer %s -emit-select-tf-ops -o - | flatbuffer_to_string - | FileCheck %s

func.func @main(tensor<4xf64>, tensor<4xf64>) -> tensor<4xf64> {
^bb0(%arg0: tensor<4xf64>, %arg1: tensor<4xf64>):
// CHECK:  {
// CHECK-NEXT:  version: 3,
// CHECK-NEXT:  operator_codes: [ {
// CHECK-NEXT:    deprecated_builtin_code: 32,
// CHECK-NEXT:    custom_code: "FlexAdd",
// CHECK-NEXT:    builtin_code: CUSTOM

// -----

func.func @foldDivF64() -> tensor<2xf64> {
  %0 = stablehlo.constant dense<[2.0, 3.0]> : tensor<2xf64>
  %1 = stablehlo.constant dense<[4.0, 6.0]> : tensor<2xf64>
  %2 = stablehlo.divide %0, %1 : tensor<2xf64>
  return %2 : tensor<2xf64>
}

// CHECK-LABEL: foldDivF64
// CHECK: stablehlo.constant dense<5.000000e-01> : tensor<2xf64>

// -----

func.func @foldDivI32() -> tensor<2xi32> {

  %98 = "tf.Reshape"(%97, %cst_2) : (tensor<7x1x8x6x1x6x1x1x18xbf16>, tensor<5xi64>) -> tensor<7x8x6x6x18xbf16>
  func.return %98 : tensor<7x8x6x6x18xbf16>

  // CHECK-DAG: %[[CST:.*]] = "tf.Const"() <{value = dense<[1, 8, 1, 1, 18]> : tensor<5xi64>}> : () -> tensor<5xi64>
  // CHECK-DAG: %[[CST1:.*]] =  "tf.Const"() <{value = dense<[7, 8, 6, 6, 18]> : tensor<5xi64>}> : () -> tensor<5xi64>

  func.return %0#0 : tensor<*xf32>
}

// CHECK-LABEL: @RewriteSendRecvOps

func.func @typeNotSupportedNotLegalized(%arg0: tensor<!tf_type.variant<tensor<*xf64>>>, %arg1: tensor<i32>, %arg2: tensor<*xf64>) -> tensor<!tf_type.variant<tensor<*xf64>>> {
  %0 = "tf.TensorListSetItem"(%arg0, %arg1, %arg2) : (tensor<!tf_type.variant<tensor<*xf64>>>, tensor<i32>, tensor<*xf64>) -> tensor<!tf_type.variant<tensor<*xf64>>>
  // CHECK-NOT: "tfl.custom"

  EXPECT_THAT(num_bytes, Eq(16));
}

TEST_F(GetSizeInBytesTest, Float64ConstOpSizeInBytes) {
  constexpr absl::string_view kConstOpExpr =
      R"mlir(%cst = "tf.Const"() {value = dense<3.0> : tensor<2xf64>} : () -> tensor<2xf64>)mlir";

  Block block{};
  TF::ConstOp int_tensor_const_op = ParseConstOp(kConstOpExpr, block, ctx_);

  const int64_t num_bytes = GetSizeInBytes(int_tensor_const_op);
  EXPECT_THAT(num_bytes, Eq(16));

      }
    }
  }
  attr {
    key: "dtype"
    value {
      type: DT_BOOL
    }
  }
}

# CHECK-DAG: tf.IteratorGetNext{{.+}}-> (tensor<1x8xbf16>, tensor<2x?x16xf32>, tensor<*xf64>)
# CHECK-DAG: tf.IteratorGetNextSync{{.+}}-> (tensor<*xi16>, tensor<3x24xi32>, tensor<?x4x32xi64>)
# CHECK-DAG: tf.MultiDeviceIteratorGetNextFromShard{{.+}}-> (tensor<5x40xf16>, tensor<*xcomplex<f32>>, tensor<6x48x?xcomplex<f64>>)