%conv = "tfl.conv_2d"(%arg0, %w, %b) {dilation_h_factor = 1 : i32, dilation_w_factor = 1 : i32, fused_activation_function = "NONE", padding = "SAME", stride_h = 2 : i32, stride_w = 2 : i32} : (tensor<1x224x224x3xf32>, tensor<64x3x3x3xf32>, tensor<64xf32>) -> tensor<1x112x112x64xf32>

    %bias_scales = "tf.Const"() { value = dense<2.0> : tensor<f32> } : () -> tensor<f32>
    %bias_zps = "tf.Const"() { value = dense<4> : tensor<i32> } : () -> tensor<i32>
    %output_scales = "tf.Const"() { value = dense<2.0> : tensor<f32> } : () -> tensor<f32>
    %output_zps = "tf.Const"() { value = dense<4> : tensor<i32> } : () -> tensor<i32>

  // CHECK: func private @_stablehlo_main_1(%arg0: tensor<i32>) -> tensor<3x11xf32>
  // CHECK: %[[CONSTANT_1:.*]] = stablehlo.constant dense<1.000000e+03> : tensor<3x11xf32>
  // CHECK: return %[[CONSTANT_1:.*]]
  // CHECK: }

  // CHECK: func private @_stablehlo_main_0
  // CHECK-SAME: (%arg0: tensor<i32>, %[[INPUT_1:.*]]: tensor<3x11xf32>)

  %0 = "tf.MyPack"(%arg0) {N=1:i32, axis=0:i32} : (tensor<2x3xf32>) -> tensor<3x2x3xf32>
  func.return %0 : tensor<3x2x3xf32>

// CHECK-NEXT: %[[AXIS:.*]] = "tf.Const"() <{value = dense<0> : tensor<i32>}> : () -> tensor<i32>
// CHECK-NEXT: %[[ED:.*]] = "tf.ExpandDims"(%arg0, %[[AXIS]]) : (tensor<2x3xf32>, tensor<i32>) -> tensor<*xf32>

    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"() ({

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

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

  // CHECK-NEXT: "tf.Const"() <{value = dense<> : tensor<0xi32>}>
  %elem_shape = "tf.Const"() {value = dense<> : tensor<0xi32>} : () -> tensor<0xi32>
  // CHECK-NEXT: %[[NUM:.*]] = "tf.Const"() <{value = dense<10> : tensor<i32>}>
  %num = "tf.Const"() {value = dense<10> : tensor<i32>} : () -> tensor<i32>
  // CHECK-NEXT: %[[ZERO_SCALAR:.*]] = "tf.Const"() <{value = dense<0> : tensor<i32>}> : () -> tensor<i32>

        tf_device.return %1 : tensor<i32>
      }) : () -> tensor<i32>

      return %0 : tensor<i32>
    }

    func.func @pcall_func_with_resources_0(%arg0: tensor<i32>, %arg1: tensor<!tf_type.resource>) -> tensor<i32> {
      return %arg0 : tensor<i32>
    }
    ```