%1 = "tfl.abs"(%arg2) {fused_activation_function = "NONE"} : (tensor<5xf32>) -> tensor<5xf32>
    %2 = stablehlo.add %1, %arg2 : tensor<5xf32>
    %3 = "tfl.abs"(%2) {fused_activation_function = "NONE"} : (tensor<5xf32>) -> tensor<5xf32>
    "stablehlo.return"(%3) : (tensor<5xf32>) -> ()
  }) {dimensions = array<i64: 0>} : (tensor<7x5xf32>, tensor<5xf32>) -> tensor<5xf32>
  func.return %0: tensor<5xf32>
}

// -----

module {
func.func @main(%arg0 : tensor<5xf32>, %arg1 : tensor<f32>, %arg2 : tensor<f32>) -> tensor<5xf32> {
  %0 = "tf.ClipByValue"(%arg0, %arg1, %arg2) : (tensor<5xf32>, tensor<f32>, tensor<f32>) -> tensor<5xf32>
  func.return %0 : tensor<5xf32>
}

// CHECK-LABEL: floor_mod
func.func @floor_mod(%arg0: tensor<5xf32>, %arg1: tensor<5xf32>) -> tensor<5xf32> {
  %0 = "tf.FloorMod"(%arg0, %arg1) : (tensor<5xf32>, tensor<5xf32>) -> tensor<5xf32>
  func.return %0 : tensor<5xf32>
// CHECK: %[[CUSTOM_0:.*]] = "tfl.custom"(%arg0, %arg1) <{custom_code = "FlexFloorMod", custom_option = #tfl<const_bytes : "{{.*}}">}> : (tensor<5xf32>, tensor<5xf32>) -> tensor<5xf32>
// CHECK: return %[[CUSTOM_0]] : tensor<5xf32>
}

// CHECK-LABEL: floor_mod
func.func @floor_mod(%arg0: tensor<5xf32>, %arg1: tensor<5xf32>) -> tensor<5xf32> {
  %0 = "tf.FloorMod"(%arg0, %arg1) : (tensor<5xf32>, tensor<5xf32>) -> tensor<5xf32>
  func.return %0 : tensor<5xf32>
// CHECK: %[[CUSTOM_0:.*]] = "tfl.custom"(%arg0, %arg1) <{custom_code = "FlexFloorMod", custom_option = #tfl<const_bytes : "{{.*}}">}> : (tensor<5xf32>, tensor<5xf32>) -> tensor<5xf32>
// CHECK: return %[[CUSTOM_0]] : tensor<5xf32>
}

    %outputs_6:2, %control_7 = tf_executor.island wraps "tf.While"(%arg0, %arg1) {body = @body, cond = @cond, is_stateless = false, shape_invariant} : (tensor<i32>, tensor<5xf32>) -> (tensor<i32>, tensor<5xf32>) loc("WhileWithOutputShapes")
    tf_executor.fetch %outputs_2#1, %outputs_4#1, %outputs_6#1 : tensor<5xf32>, tensor<5xf32>, tensor<5xf32>
  }

  %1 = "tf.Mul"(%arg1, %arg1) : (tensor<2xf32>, tensor<2xf32>) -> tensor<2xf32>
  %2 = "tf.Div"(%arg1, %arg1) : (tensor<2xf32>, tensor<2xf32>) -> tensor<2xf32>

  // CHECK: "tf.If"{{.+}}else_branch = @tf.IfRegion_else{{.+}}then_branch = @tf.IfRegion_then
  %3 = "tf.IfRegion"(%arg0) ({
     %4 = "tf.Select"(%arg0, %0, %1) : (tensor<i1>, tensor<2xf32>, tensor<2xf32>) -> tensor<2xf32>
     "tf.Yield"(%4) : (tensor<2xf32>) -> ()

// CHECK-LABEL: func @init
func.func @init(%arg0: tensor<4xf32>) -> tensor<7xf32> {
    %0 = builtin.unrealized_conversion_cast to tensor<7xf32>
    return %0 : tensor<7xf32>
}

// CHECK-LABEL: func @next
func.func @next(%arg0: tensor<7xf32>, %arg1: tensor<3xf32>) -> tensor<6xf32> {
    %0 = builtin.unrealized_conversion_cast to tensor<6xf32>
    return %0 : tensor<6xf32>
}

// CHECK-LABEL: func @finalize

    func.return %1, %0 : tensor<?x5xf32>, tensor<?x5xf32>
  }

  func.func @main() -> (tensor<5xf32> {tf_saved_model.index_path = ["start_logits"]}) attributes {tf.entry_function = {control_outputs = "", inputs = "", outputs = "StatefulPartitionedCall:1"}, tf_saved_model.exported_names = ["serving_default"]} {
    %cst = arith.constant dense<0.000000e+00> : tensor<5xf32>
    func.return %cst : tensor<5xf32>
  }

func.func @squeeze_folder(%arg0 : tensor<?x?xf32>) -> tensor<?x?xf32> {
  %0 = "tfl.squeeze"(%arg0) : (tensor<?x?xf32>) -> tensor<?x?xf32>
  // CHECK: return %arg0
  func.return %0 : tensor<?x?xf32>
}

// -----

// CHECK-LABEL: @squeeze_folder
func.func @squeeze_folder(%arg0 : tensor<?x?xf32>) -> tensor<*xf32> {
  %0 = "tfl.squeeze"(%arg0) : (tensor<?x?xf32>) -> tensor<*xf32>
  // CHECK: "tfl.squeeze"