%cst = arith.constant dense<[3, 5, 7]> : tensor<3xi32>
  %0 = "tf.BroadcastTo"(%arg1, %cst) : (tensor<5xf32>, tensor<3xi32>) -> tensor<3x5x7xf32>
  %1 = "tf.Mul"(%arg0, %0) : (tensor<5x7xf32>, tensor<3x5x7xf32>) -> tensor<3x5x7xf32>
  func.return %1 : tensor<3x5x7xf32>
  // CHECK: %[[C0:.*]] = arith.constant dense<[3, 5, 7]> : tensor<3xi32>

func.func @op_with_tfl_control_flow() -> (tensor<1xf32>, !tfl.control) {
  // CHECK: vhlo.constant_v1
  %0 = stablehlo.constant dense<1.000000e+00> : tensor<1xf32>
  // CHECK-NEXT: tfl.control_node
  %outputs, %control = tfl.control_node {
    %1 = "tfl.neg"(%0) : (tensor<1xf32>) -> tensor<1xf32>
    "tfl.yield"(%1) : (tensor<1xf32>) -> ()
  }
  return %outputs, %control : tensor<1xf32>, !tfl.control
}

// -----

  func.func @init() attributes {tf_saved_model.exported_names = ["__tf_saved_model_session_initializer"], tf_saved_model.initializer_type = "restore_op"} {
    "tf.Const"() {value = dense<[1.0]> : tensor<1xf32> } : () -> tensor<1xf32>
    func.return
  }
}

// -----

module attributes {tf_saved_model.semantics, tf_saved_model.under_construction} {
  // Test case: Variable removal.
  // CHECK: func @init()

// %0 = "tfl.transpose"(%arg0, %cst) : (tensor<56x8x56x1x1x1x7xf32>, tensor<7xi32>) -> tensor<1x1x8x56x56x7x1xf32>
// to-
// %0 = "tfl.reshape"(%arg0, %cst) : (tensor<56x8x56x1x1x1x7xf32>, tensor<4xi32>) -> tensor<56x8x56x7xf32>
// %1 = "tfl.transpose"(%0, %cst_0) : (tensor<56x8x56x7xf32>, tensor<4xi32>) -> tensor<8x56x56x7xf32>
// %2 = "tfl.reshape"(%1, %cst_1) : (tensor<8x56x56x7xf32>, tensor<7xi32>) -> tensor<1x1x8x56x56x7x1xf32>

// Confirm function references in if ops are preserved
func.func @main(%arg0: tensor<1xf32>, %arg1: tensor<1xf32>) -> tensor<1xf32> {
// CHECK:   %{{.*}} = "tf.If"(%{{.*}}, %{{.*}}, %{{.*}}) <{else_branch = @cond_false, is_stateless = false, then_branch = @cond_true}> : (tensor<1xi1>, tensor<1xf32>, tensor<1xf32>) -> tensor<1xf32>
  %0 = "tfl.less"(%arg0, %arg1) : (tensor<1xf32>, tensor<1xf32>) -> tensor<1xi1>

module {
func.func @main(%arg0: tensor<2xi32>) -> tensor<2xi32> attributes {tf.entry_function = {inputs = "arg0", outputs = "tfl.custom1"}} {
  %0 = "tfl.custom"(%arg0, %arg0) {custom_code = "stablehlo.subtract", custom_option = #tfl<const_bytes : "0x00000100002401">} : (tensor<2xi32>, tensor<2xi32>) -> tensor<2xi32>
  func.return %0 : tensor<2xi32>
}
}

// CHECK:       module {

    %4:2 = tf_executor.island wraps "tf.Add"(%arg0, %3) {T = "tfdtype$DT_INT32", device = ""} : (tensor<10xi32>, tensor<10xi32>) -> tensor<10xi32> loc("Add")
    tf_executor.fetch %4#0 : tensor<10xi32>
  }
  func.return %graph : tensor<10xi32>
}

// CHECK: name: "foo"
// CHECK-NEXT: op: "_Arg"
// CHECK: name: "bar"

      "tfl.yield"(%wrap_result): (tensor<1xf32>)->()
      }): (tensor<1xf32>, tensor<1xf32>) -> (tensor<1xf32>)
  %tmp3 = "tfl.add"(%const, %tmp2) { fused_activation_function = "NONE" } : (tensor<1xf32>, tensor<1xf32>) -> (tensor<1xf32>)
  %tmp4 = "tf.AddV2"(%tmp2, %tmp2) { device = "" } : (tensor<1xf32>, tensor<1xf32>) -> (tensor<1xf32>)

// Inline a function that contains only tfl ops.
func.func @func_with_tfl_ops(%arg0 : tensor<2xi32>) -> tensor<2xi32> {
  %0 = "tfl.sub"(%arg0, %arg0) {fused_activation_function = "RELU6"} : (tensor<2xi32>, tensor<2xi32>) -> tensor<2xi32>
  %1 = "tfl.add"(%0, %arg0) {fused_activation_function = "RELU6"} : (tensor<2xi32>, tensor<2xi32>) -> tensor<2xi32>
  func.return %1: tensor<2xi32>
}

// CHECK-LABEL: func @inline_with_arg

    return
  }
  // CHECK-LABEL: func private @foo.impl
  func.func private @foo.impl(%arg0: tensor<1xf32>, %arg1: tensor<2xf32>) -> (tensor<2xf32>, tensor<1xf32>) {
    return %arg1, %arg0 : tensor<2xf32>, tensor<1xf32>
  }