tf_device.return %3 : tensor<i32>
  }) {device = "/device:TPU:0"} : () -> tensor<i32>
  func.return %execute : tensor<i32>
}

// -----

// Tests that the pass does not change unsupported input ops.

// CHECK-LABEL: func @unsupported_ops
func.func @unsupported_ops(%arg0: tensor<3x3x1x32xf32> {tf.device = "/device:CPU:0"}) -> tensor<i32> {
  %compile:2 = "tf_device.launch"() ({

  %0 = "tfl.while"(%arg0) ({
  ^bb0(%block: tensor<i32>):
    %1 = "tfl.add"(%arg0, %arg0) { fused_activation_function = "NONE", tac.device = "CPU", tac.inference_type = "FLOAT"} : (tensor<i32>, tensor<i32>) -> (tensor<i32>)
    %2 = "tfl.add"(%1, %1) { fused_activation_function = "NONE", tac.device = "GPU", tac.inference_type = "FLOAT"} : (tensor<i32>, tensor<i32>) -> (tensor<i32>)
    "tfl.yield"(%2) : (tensor<i32>) -> ()
  },{

    %Const, %ctl = Const name("Const") {dtype = i32, value = dense<0> : tensor<i32>} : () -> (tensor<i32>)
    %Enter, %ctl_0 = Enter(%Const) name("while/Enter") {T = i32, frame_name = "while/while_context", is_constant = false, parallel_iterations = 10 : i64} : (tensor<i32>) -> (tensor<*xi32>)
    %NextIteration, %ctl_1 = NextIteration(%Add) name("while/NextIteration") {T = i32} : (tensor<*xi32>) -> (tensor<*xi32>)

  }) : (tensor<i32>, tensor<1xf32>, tensor<i32>) -> (tensor<i32>, tensor<1xf32>, tensor<i32>)
  func.return %0#1, %0#2 : tensor<1xf32>, tensor<i32>

    %0:2 = "tf_device.launch"() ({
	  %1 = "tf.Cast"(%arg0) {Truncate = false} : (tensor<2048xi64>) -> tensor<2048xi32>
	  %2 = "tf.Cast"(%arg1) {Truncate = false} : (tensor<2048xi64>) -> tensor<2048xi32>

  }) {n = 2 : i32, operandSegmentSizes = array<i32: 0, 0>} : () -> ()
  func.return
}

// -----

// Check that an empty replicate block is invalid.
func.func @verifier_replicate_empty_block() {
  "tf_device.replicate" () ({
// expected-error@-1 {{'tf_device.replicate' op expects a non-empty block}}
  ^entry:
  }) {n = 2 : i32, operandSegmentSizes = array<i32: 0, 0>} : () -> ()
  func.return
}

// CHECK: func @main([[MAIN_ARG0:%.*]]: tensor<i32>) -> tensor<i32>
func.func @main(%arg0: tensor<i32>) -> tensor<i32> {
  // CHECK:      [[MAIN_TOKEN:%.*]] = mhlo.create_token

  // CHECK:      [[CALL:%.*]]:2 = call @callee([[MAIN_ARG0]], [[MAIN_TOKEN]])
  // CHECK-SAME: (tensor<i32>, !mhlo.token) -> (tensor<i32>, !mhlo.token)
  %0 = func.call @callee(%arg0) : (tensor<i32>) -> tensor<i32>

  // CHECK:      return [[CALL]]#0

    // expected-remark@+1 {{Cost: 1}}
    %12 = "tf.StridedSlice"(%8, %6, %1, %4) {begin_mask = 0 : i64, device = "/job:localhost/replica:0/task:0/device:CPU:0", ellipsis_mask = 0 : i64, end_mask = 0 : i64, new_axis_mask = 0 : i64, shrink_axis_mask = 1 : i64} : (tensor<4xi32>, tensor<1xi32>, tensor<1xi32>, tensor<1xi32>) -> tensor<i32>
    // expected-remark@+1 {{Cost: 5}}

// CHECK-LABEL: import_stats_skip
func.func @import_stats_skip(%arg0: tensor<4xf32>, %cst: tensor<i32>) -> (tensor<2xf32>,tensor<2xf32>) {
  %0:2 = "tfl.split"(%cst, %arg0) {num_splits = 2 : i32} : (tensor<i32>, tensor<4xf32>) -> (tensor<2xf32>, tensor<2xf32>)
    loc(fused["skip1", "skip2.cc":10:8, callsite("op" at "skip3.cc":10:8)])
  func.return %0#0, %0#1 : tensor<2xf32>, tensor<2xf32>

    %2 = "tf.Const"() {value = dense<1> : tensor<i32>} : () -> tensor<i32>
    %3 = "tf.Const"() {value = dense<1> : tensor<i32>} : () -> tensor<i32>
    %4 = "tf.Add"(%2, %3) : (tensor<i32>, tensor<i32>) -> tensor<i32>
    "tf.AssignAddVariableOp"(%0, %4) {dtype = "tfdtype$DT_INT32"} : (tensor<*x!tf_type.resource<tensor<2x8xi32>>>, tensor<i32>) -> ()