%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
}

// -----

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

    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>
  }

func.func private @simple_callee() -> tensor<2xi32>  {
  %cst = "tf.Const"() { value = dense<2> : tensor<2xi32> } : () -> tensor<2xi32>
  func.return %cst : tensor<2xi32>
}

func.func private @simple_callee_with_noinline() -> tensor<2xi32> attributes {tf._noinline = true} {
  %cst = "tf.Const"() { value = dense<2> : tensor<2xi32> } : () -> tensor<2xi32>
  func.return %cst : tensor<2xi32>
}

// Test that simple TF operations can be inlined.

    // CHECK:     return %[[RESULT]]
    func.return %0 : tensor<10xi32>
  }

  // CHECK-LABEL: func @foo
  // CHECK-SAME:    %[[ARG0:.*]]: tensor<10xi32>, %[[ARG1:.*]]: tensor<10xi32>
  func.func @foo(%arg0: tensor<10xi32>, %arg1: tensor<10xi32>) -> tensor<10xi32> {

    // CHECK: call @f
    %arg2 = func.call @f(%arg0) : (tensor<?xi32>) -> (tensor<?xi32>)
    return %arg2 : tensor<?xi32>
  }

  // CHECK-LABEL: func private @f
  // CHECK:    attributes {_from_xla_call_module}
  func.func private @f(%arg0: tensor<?xi32>) -> (tensor<?xi32>) attributes {_from_xla_call_module} {
    return %arg0 : tensor<?xi32>

    } : (tensor<10xf32>, tensor<10xf32>, tensor<6xi32>) -> (tensor<10xf32>, tensor<6xi32>)
    // CHECK: return %[[RESULT_CAST]], %[[RESULT]]#1 : tensor<10xf32>, tensor<6xi32>
    func.return %0#0, %0#1 : tensor<10xf32>, tensor<6xi32>
  }

  // CHECK-LABEL: func private @main_0
  // CHECK-SAME: %[[ARG_0:.*]]: tensor<10xbf16>, %[[ARG_1:.*]]: tensor<10xbf16>, %[[ARG_2:.*]]: tensor<6xi32>
  func.func private @main_0(

  %fill = "tf.Fill"(%cst_0, %cst) : (tensor<1xi64>, tensor<f32>) -> tensor<1xf32>
  tf_device.replicate([%0, %fill] as %arg_r0: tensor<1xf32>) {n = 2 : i32} {
    %1 = "tf_device.launch"() <{device = "TPU_REPLICATED_HOST_0"}> ({
      %2 = "tf.Identity"(%arg_r0) : (tensor<1xf32>) -> tensor<1xf32>
      tf_device.return %2 : tensor<1xf32>
    }) : () -> tensor<1xf32>