%ch0 = tfrt.new.chain

  %a1 = tfrt_fallback_async.const_dense_tensor dense<[[2, 2], [2, 2]]> : tensor<2x2xi32>
  %a2 = tfrt_fallback_async.const_dense_tensor dense<[[3, 3], [3, 3]]> : tensor<2x2xi32>
  %b = tfrt_fallback_async.const_dense_tensor dense<[[1, 1], [1, 1]]> : tensor<2x2xi32>

  // Two batch_size=2 batches get concatenated.

    %12 = "tfl.logistic"(%11) : (tensor<4x2xf32>) -> tensor<4x2xf32>
    %13 = tfl.mul %arg4, %12 {fused_activation_function = "NONE"} : tensor<4x2xf32>
    %14 = "tfl.relu"(%10#1) : (tensor<4x2xf32>) -> tensor<4x2xf32>
    %15 = "tfl.logistic"(%10#0) : (tensor<4x2xf32>) -> tensor<4x2xf32>
    %16 = tfl.mul %15, %14 {fused_activation_function = "NONE"} : tensor<4x2xf32>

// CHECK-LABEL: main
func.func @main() -> tensor<3x2xi32> {
  // CHECK: "tfl.pseudo_const"() <{value = dense<0> : tensor<3x2xi32>}> {tfl.is_variable} : () -> tensor<3x2xi32>
  %0 = "tfl.pseudo_const"() {value = dense<0> : tensor<3x2xi32>, tfl.is_variable} : () -> tensor<3x2xi32> loc("variable")
  func.return %0 : tensor<3x2xi32>

module {
  func.func @main(%arg0: tensor<1x2xi32>) -> tensor<1x2x2xi32> {
  %0 = "tfl.custom"(%arg0) {custom_code = "stablehlo.broadcast_in_dim", custom_option = #tfl<const_bytes : "0x62726F6164636173745F64696D656E73696F6E73000201020119010101072C022401">} : (tensor<1x2xi32>) -> tensor<1x2x2xi32>
  func.return %0 : tensor<1x2x2xi32>
  }
}

// CHECK:      module {
// CHECK-NEXT:  func @main(%arg0: tensor<1x2xi32>) -> tensor<1x2x2xi32> {

  %1 = "tfl.sub" (%arg0, %0) {fused_activation_function = "RELU6"} : (tensor<3x2xi32>, tensor<3x2xi32>) -> tensor<3x2xi32> loc("sub")
  %2 = "arith.constant" () {value = dense<10> : tensor<i32>} : () -> tensor<i32> loc("SameNameAsOutput")
  %3 = "tfl.add" (%2, %1) {fused_activation_function = "NONE"} : (tensor<i32>, tensor<3x2xi32>) -> tensor<3x2xi32> loc("add")

// CHECK-NEXT: }

  %0 = "tfl.pseudo_const" () {value = dense<[[1, 2], [3, 4], [5, 6]]> : tensor<3x2xi32>} : () -> tensor<3x2xi32> loc("Const")
  %1,%control_1 = tfl.control_node controls "tfl.sub" (%arg0, %0) {fused_activation_function = "RELU6"} : (tensor<3x2xi32>, tensor<3x2xi32>) -> tensor<3x2xi32> loc("sub")
  %2 = "arith.constant" () {value = dense<10> : tensor<i32>} : () -> tensor<i32> loc("SameNameAsOutput")

// RUN: flatbuffer_translate -mlir-to-tflite-flatbuffer %s -o - | flatbuffer_to_string - | FileCheck %s

func.func @main() -> tensor<3x2xi32> {
  %0 = "tfl.pseudo_const" () {value = dense<0> : tensor<3x2xi32>, tfl.is_variable} : () -> tensor<3x2xi32> loc("variable")
  func.return %0 : tensor<3x2xi32>
}

// CHECK:      {
// CHECK-NEXT:     version: 3,
// CHECK-NEXT:     operator_codes: [ ],
// CHECK-NEXT:     subgraphs: [ {
// CHECK-NEXT:       tensors: [ {

  %2:2 = "tf.RecvTPUEmbeddingActivations"() {_tpu_embedding_layer = "call1", config = "\0A\0B\0C\0D"} : () -> (tensor<2x2xf32>, tensor<4x4xf32>)
  func.return
}

// -----

// RUN: flatbuffer_translate -mlir-to-tflite-flatbuffer %s -o - | flatbuffer_to_string - | FileCheck %s

func.func @main(tensor<3x2xi32>) -> tensor<6xi32> {
^bb0(%arg0: tensor<3x2xi32>):
// CHECK: {
// CHECK-NEXT:   version: 3,
// CHECK-NEXT:   operator_codes: [ {
// CHECK-NEXT:     deprecated_builtin_code: 22,
// CHECK-NEXT:     version: 1,
// CHECK-NEXT:     builtin_code: RESHAPE
// CHECK-NEXT:   } ],
// CHECK-NEXT:   subgraphs: [ {

  %1 = tfl.add %0, %cst { fused_activation_function = "NONE" } : tensor<3x2xi32>
  func.return %1 : tensor<3x2xi32>
}

// CHECK: %cst = arith.constant dense<[1, 0]> : tensor<2xi32>
// CHECK: [1, 3, 5], [2, 4, 6]
// CHECK: %0 = tfl.add %arg0, %cst_0 {fused_activation_function = "NONE"} : tensor<2x3xi32>
// CHECK: %1 = "tfl.transpose"(%0, %cst) : (tensor<2x3xi32>, tensor<2xi32>) -> tensor<3x2xi32>
// CHECK: return %1 : tensor<3x2xi32>