// CHECK: while_1([[ARG0:%[^ :]*]]: tensor<i32>, [[ARG1:%[^ :]*]]: tensor<1xf32>)
func.func @while_1(%arg0: tensor<i32>, %arg1: tensor<1xf32>) -> tensor<1xf32> {
  // CHECK: [[CST:%[^ ]*]] = arith.constant dense<1> : tensor<i32>
  // CHECK: "tfl.while"([[ARG0]], [[ARG1]])
  // CHECK: (tensor<i32>, tensor<1xf32>) -> (tensor<i32>, tensor<1xf32>)
  %0:2 = "tfl.while"(%arg0, %arg1) (
    // cond
    {

    %a:2, %control1 = tf_executor.island wraps "tf.A"() {_class = ["loc:@class"]} : () -> (tensor<2xf32>, tensor<2xf32>)
    %s:2, %control2 = tf_executor.island wraps "tf.Split"(%c, %a#1) {num_split = 2 : i32} : (tensor<i32>, tensor<2xf32>) -> (tensor<1xf32>, tensor<1xf32>)
    tf_executor.fetch
  }
  func.return
}

// -----

// CHECK-LABEL: func @no_colocate_split_has_device

  ^bb1(%0: tensor<!tf_type.resource<tensor<?xf32>>>, %1: tensor<!tf_type.resource<tensor<?xf32>>>, %2: tensor<!tf_type.resource<tensor<?xf32>>>):
    "tf.AssignVariableOp"(%0, %arg0) : (tensor<!tf_type.resource<tensor<?xf32>>>, tensor<?xf32>) -> ()
    cf.br ^bb1(%1, %2, %0 : tensor<!tf_type.resource<tensor<?xf32>>>, tensor<!tf_type.resource<tensor<?xf32>>>, tensor<!tf_type.resource<tensor<?xf32>>>)
  }

// CHECK:       "tfl.yield"(%{{.*}}) : (tensor<i1>) -> ()
// CHECK:     },  {
// CHECK:     ^bb0(%{{.*}}: tensor<*xi32>, %{{.*}}: tensor<*xf32>):
// CHECK:       "tfl.yield"(%{{.*}}, %{{.*}}) : (tensor<*xi32>, tensor<*xf32>) -> ()
// CHECK:     }) : (tensor<i32>, tensor<1xf32>) -> (tensor<i32>, tensor<1xf32>)

func.func @main(%arg0: tensor<i32>, %arg1: tensor<1xf32>) -> tensor<1xf32> {

  %min = arith.constant dense<0.0> : tensor<16xf32>
  %max = arith.constant dense<255.0> : tensor<16xf32>
  %mini = "tf.Identity"(%min) : (tensor<16xf32>) -> tensor<16xf32>
  %maxi = "tf.Identity"(%max) : (tensor<16xf32>) -> tensor<16xf32>
  %fq = "tf.FakeQuantWithMinMaxVarsPerChannel"(%in, %mini, %maxi) {num_bits = 5, narrow_range = false} : (tensor<3x3x3x16xf32>, tensor<16xf32>, tensor<16xf32>) -> tensor<3x3x3x16xf32>

  // expected-error@+1 {{'mhlo.uniform_dequantize' op is illegal as it is a UQ op or contains uq/qint types}}
  %0 = mhlo.uniform_dequantize %arg0 : (tensor<?x?x!quant.uniform<i8:f32, 1.000000e+00:3>>) -> tensor<?x?xf32>
  return %0 : tensor<?x?xf32>
}

// -----

func.func @illegal_tf_uniform_quantize(%arg0 : tensor<1xf32>) -> tensor<1xf32> {

        %6:2 = tf_executor.island wraps "tf.opC"() : () -> tensor<1xf32>
        %8:2 = tf_executor.island(%4#1) wraps "tf.opD"(%6#0) : (tensor<1xf32>) -> tensor<1xf32>
        tf_executor.fetch %8#0 : tensor<1xf32>
      }
      tf_executor.yield %3 : tensor<1xf32>
    }
    tf_executor.fetch
  }
  func.return
}

// CHECK:        tf_executor.island wraps "tf.opA"

  %min = arith.constant dense<0.0> : tensor<16xf32>
  %max = arith.constant dense<15.0> : tensor<16xf32>
  %mini = "tf.Identity"(%min) : (tensor<16xf32>) -> tensor<16xf32>
  %maxi = "tf.Identity"(%max) : (tensor<16xf32>) -> tensor<16xf32>
  %fq = "tf.FakeQuantWithMinMaxVarsPerChannel"(%in, %mini, %maxi) {num_bits = 3, narrow_range = false} : (tensor<3x3x3x16xf32>, tensor<16xf32>, tensor<16xf32>) -> tensor<3x3x3x16xf32>

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

  },  {
    ^bb0(%arg2: tensor<*xi32>, %arg3: tensor<*xf32>, %arg4: tensor<i32>):
      %1 = "tfl.sub"(%arg2, %arg4) {fused_activation_function = "NONE"} :
        (tensor<*xi32>, tensor<i32>) -> tensor<*xi32>
      %2 = tfl.add %arg3, %arg3 {fused_activation_function = "NONE"} : tensor<*xf32>
      "tfl.yield"(%1, %2, %arg4) : (tensor<*xi32>, tensor<*xf32>, tensor<i32>) -> ()
  }) : (tensor<i32>, tensor<1xf32>, tensor<i32>) -> (tensor<i32>, tensor<1xf32>, tensor<i32>)