// CHECK-DAG:  %[[CST0:.*]] = "tf.Const"() <{value = dense<[3, 2, 4]> : tensor<3xi64>}> : () -> tensor<3xi64>
  // CHECK-DAG:  %[[CST1:.*]] = "tf.Const"() <{value = dense<0> : tensor<3xi64>}> : () -> tensor<3xi64>
  // CHECK-DAG:  %[[CST2:.*]] = "tf.Const"() <{value = dense<[3, 6, 4]> : tensor<3xi64>}> : () -> tensor<3xi64>

  %cst0 = mhlo.constant dense<[1, 2, 3, 4]> : tensor<4xi32>
  // CHECK: %[[BROADCAST:.*]] = "mhlo.broadcast_in_dim"(%[[CONST]]) <{broadcast_dimensions = dense<3> : tensor<1xi64>}> : (tensor<4xi32>) -> tensor<1x1x2x4xi32>
  %0 = "mhlo.broadcast_in_dim"(%cst0) <{broadcast_dimensions = dense<3> : tensor<1xi64>}> : (tensor<4xi32>) -> tensor<1x1x2x4xi32>
  // CHECK: %[[MUL:.*]] = mhlo.multiply %[[BROADCAST]], %[[ARG]] : tensor<1x1x2x4xi32>

  %cst_0 = arith.constant dense<[2, -1]> : tensor<2xi64>
  %0 = "tf.FakeQuantWithMinMaxVars"(%arg0, %arg1, %arg2) {num_bits = 5, narrow_range = false} : (tensor<1x2xf32>, tensor<f32>, tensor<f32>) -> tensor<1x2xf32>
  %1 = "tf.Reshape"(%0, %cst_0) : (tensor<1x2xf32>, tensor<2xi64>) -> tensor<2x1xf32>
  %2 = "tf.Reshape"(%0, %cst_0) : (tensor<1x2xf32>, tensor<2xi64>) -> tensor<2x1xf32>
  func.return %1, %2 : tensor<2x1xf32>, tensor<2x1xf32>

      // CHECK:      name: "ParseExample"
      // CHECK-NEXT: op: "ParseExampleV2"
      // CHECK-NEXT: input: "input0"

  %cst_0 = arith.constant dense<[2, -1]> : tensor<2xi64>
  %0 = "tf.FakeQuantWithMinMaxVars"(%arg0, %arg1, %arg2) {num_bits = 3, narrow_range = false} : (tensor<1x2xf32>, tensor<f32>, tensor<f32>) -> tensor<1x2xf32>
  %1 = "tf.Reshape"(%0, %cst_0) : (tensor<1x2xf32>, tensor<2xi64>) -> tensor<2x1xf32>
  %2 = "tf.Reshape"(%0, %cst_0) : (tensor<1x2xf32>, tensor<2xi64>) -> tensor<2x1xf32>
  func.return %1, %2 : tensor<2x1xf32>, tensor<2x1xf32>

    -> (tensor<?x2xi64>, tensor<?x2xi64>, tensor<?x!tf_type.string>, tensor<?xi64>, tensor<2xi64>, tensor<2xi64>, tensor<32xf32>, tensor<32xf32>)

  // CHECK: tfrt.return {{.*}}, [[outputs]]#0
  func.return %outputs#0 : tensor<?x2xi64>
}

// CHECK-LABEL: func @no_native

    %24 = "tf.Softmax"(%18) : (tensor<2x1001xf32>) -> tensor<2x1001xf32>
    %25 = "tf.ArgMax"(%24, %2) : (tensor<2x1001xf32>, tensor<i32>) -> tensor<2xi64>
    %26 = "tf.Cast"(%25) {Truncate = false} : (tensor<2xi64>) -> tensor<2xf32>
    %27 = "tf.Equal"(%14, %26) {incompatible_shape_error = true} : (tensor<2xf32>, tensor<2xf32>) -> tensor<2xi1>

    // CHECK: "mhlo.pad"(%[[ARG0]], %[[ARG1]])
    // CHECK-SAME-DAG: edge_padding_high = dense<[1, 2]> : tensor<2xi64>
    // CHECK-SAME-DAG: edge_padding_low = dense<[2, 1]> : tensor<2xi64>
    // CHECK-SAME-DAG: interior_padding = dense<[1, 0]> : tensor<2xi64>

    %0 = mhlo.constant dense<[2, 1]> : tensor<2xi32>
    %1 = mhlo.constant dense<[1, 2]> : tensor<2xi32>

  %2 = "tf.Shape"(%0) {device = "/CPU:0"} : (tensor<*xi32>) -> tensor<?xi64>
  %dyn_2 = "tf.SetStaticDimensionBounds" (%2, %arg4) :(tensor<?xi64>, tensor<?xi64>) -> tensor<?xi64>
  %compilation_status, %program = "tf._TPUCompileMlir"(%1, %2) {device = "/CPU:0", metadata = "metadata", mlir_module = "mlir_module"} : (tensor<?xi64>, tensor<?xi64>) -> (tensor<!tf_type.string>, tensor<3x!tf_type.string>)

func.func private @while_body(%arg0 : tensor<4xf32>, %arg1 : tensor<i32>) -> (tensor<4xi64>, tensor<i32>)
func.func @testWhileRegionIncompatibleCast(%arg0 : tensor<*xi64>, %arg1 : tensor<i32>) -> tensor<*xi64> {
  // CHECK: [[Result:%.*]]:2 = "tf.While"(%arg0, %arg1) <{body = @tf.WhileRegion_body, cond = @tf.WhileRegion_cond
  %0:2 = "tf.WhileRegion"(%arg0, %arg1) (
    {
      ^bb0(%carg0: tensor<*xi64>, %carg1: tensor<i32>):