"tfl.yield"(%1#0, %1#1) : (tensor<*xi32>, tensor<*xf32>) -> ()
  }) : (tensor<i32>, tensor<1xf32>) -> (tensor<i32>, tensor<1xf32>) loc("WhileOp")
  // CHECK: (tensor<i32>, tensor<1xf32>, tensor<i32>) ->
  // CHECK-SAME: (tensor<i32>, tensor<1xf32>, tensor<i32>)
  func.return %0#1 : tensor<1xf32>
}

func.func private @WhileOp_cond(%arg0: tensor<*xi32>, %arg1: tensor<*xf32>, %arg2: tensor<i32>) -> tensor<i1> {

  } : (tensor<1x!tf_type.qint8>, tensor<f32>, tensor<i32>) -> tensor<1xf32>

  // CHECK: return %[[y]] : tensor<1xf32>
  func.return %0 : tensor<1xf32>
}

// -----

// CHECK-LABEL: func @uniform_quantize_dequantize
func.func @uniform_quantize_dequantize(%arg0: tensor<1xf32>) -> tensor<1xf32>
{
  %scales = "tf.Const"() { value = dense<1.0> : tensor<f32> } : () -> tensor<f32>

func.func @main(
    %arg0: tensor<10x10xf32>, %scale: tensor<10xf32>, %zp: tensor<10xi32>
  ) -> tensor<10x10xi8> {
  %0 = "tf.UniformQuantize"(%arg0, %scale, %zp) {
    quantization_axis = 1 : i64,
    quantization_min_val = -128 : i64,
    quantization_max_val = 127 : i64
  } : (tensor<10x10xf32>, tensor<10xf32>, tensor<10xi32>) -> tensor<10x10x!tf_type.qint8>

  func.return %0 : tensor<*xf32>
}

// CHECK-NOT: "tfl.reshape"
// CHECK: "tfl.pack"

// -----

func.func @squaredDifference(%arg0: tensor<4xf32>, %arg1: tensor<4xf32>) -> tensor<4xf32> {
  %0 = "tfl.squared_difference"(%arg0, %arg1) : (tensor<4xf32>, tensor<4xf32>) -> tensor<4xf32>
  func.return %0 : tensor<4xf32>
}

    %4 = "tf.BiasAdd"(%3, %cst) {data_format = "NHWC", device = ""} : (tensor<*xf32>, tensor<2xf32>) -> tensor<*xf32>
    %5 = "tf.Relu6"(%4) {device = ""} : (tensor<*xf32>) -> tensor<*xf32>

    func.return %2, %5 : tensor<*xf32>, tensor<*xf32>
  }

  func.func @float_conv_1(%arg0: tensor<1x3x4x3xf32>, %arg1: tensor<2x3x3x2xf32>) -> (tensor<*xf32>, tensor<*xf32>) {

  %0 = "quantfork.stats"(%arg0) {layerStats = dense<[-1.0, 1.0]> : tensor<2xf32>} : (tensor<1x1xf32>) -> tensor<1x1xf32>
  %1 = "tfl.tanh"(%0) : (tensor<1x1xf32>) -> tensor<1x1xf32>
  %2 = "quantfork.stats"(%1) {layerStats = dense<[-1.0, 1.0]> : tensor<2xf32>} : (tensor<1x1xf32>) -> tensor<1x1xf32>
  func.return %2 : tensor<1x1xf32>

  func.return %0 : tensor<? x ? x ? x f32>
}

// CHECK-LABEL: func @init
func.func @init(%arg0: tensor<4xf32>) -> tensor<7xf32> {
    %0 = builtin.unrealized_conversion_cast to tensor<7xf32>
    return %0 : tensor<7xf32>
}

// CHECK-LABEL: func @next
func.func @next(%arg0: tensor<7xf32>, %arg1: tensor<3xf32>) -> tensor<6xf32> {
    %0 = builtin.unrealized_conversion_cast to tensor<6xf32>
    return %0 : tensor<6xf32>
}

// CHECK-LABEL: func @finalize

  %1 = "tf.Mul"(%arg1, %arg1) : (tensor<2xf32>, tensor<2xf32>) -> tensor<2xf32>
  %2 = "tf.Div"(%arg1, %arg1) : (tensor<2xf32>, tensor<2xf32>) -> tensor<2xf32>

  // CHECK: "tf.If"{{.+}}else_branch = @tf.IfRegion_else{{.+}}then_branch = @tf.IfRegion_then
  %3 = "tf.IfRegion"(%arg0) ({
     %4 = "tf.Select"(%arg0, %0, %1) : (tensor<i1>, tensor<2xf32>, tensor<2xf32>) -> tensor<2xf32>
     "tf.Yield"(%4) : (tensor<2xf32>) -> ()

// PerChannel-DAG: %[[scale2:.*]] = "tf.Const"() <{value = dense<{{[0-9\.Ee\+\-]+}}> : tensor<6xf32>}> : () -> tensor<6xf32>
// PerChannel-DAG: %[[zp1:.*]] = "tf.Const"() <{value = dense<{{[0-9]+}}> : tensor<3xi32>}> : () -> tensor<3xi32>
// PerChannel-DAG: %[[zp2:.*]] = "tf.Const"() <{value = dense<{{[0-9]+}}> : tensor<6xi32>}> : () -> tensor<6xi32>