%b = builtin.unrealized_conversion_cast to tensor<f32>
    %0, %1, %2 = "tf.WhileRegion"(%a, %b, %arg0) ({
      ^bb0(%ca: tensor<f32>, %cb: tensor<f32>, %cc: !res):
      %cond = builtin.unrealized_conversion_cast to tensor<i1>
      "tf.Yield"(%cond, %ca, %cb, %cc) : (tensor<i1>, tensor<f32>, tensor<f32>, !res) -> ()
    }, {
      ^bb0(%ba: tensor<f32>, %bb: tensor<f32>, %bc: !res):

        %cast = "tf.Cast"(%output) : (tensor<1024x24x24x3xi32>) -> tensor<1024x24x24x3xf32>
        %identity = "tf.Identity"(%scale) : (tensor<f32>) -> tensor<f32>
        %mul = "tf.Mul"(%cast, %identity) : (tensor<1024x24x24x3xf32>, tensor<f32>) -> tensor<1024x24x24x3xf32>
        func.return %mul : tensor<1024x24x24x3xf32>
      }
    }
  )mlir";

}

// -----

// CHECK-LABEL: constant
func.func @constant() -> (!tfr.attr, !tfr.attr, !tfr.attr, !tfr.attr) {
  %0 = tfr.constant f32 -> !tfr.attr
  %1 = tfr.constant [f32, i32] -> !tfr.attr
  %2 = "tfr.constant"() {value = f32} : () -> !tfr.attr
  %3 = "tfr.constant"() {value = [f32, i32]} : () -> !tfr.attr
  func.return %0, %1, %2, %3 : !tfr.attr, !tfr.attr, !tfr.attr, !tfr.attr
}

// -----

// CHECK-LABEL: equal

    %round = "tf.Floor"(%add) : (tensor<*xf32>) -> tensor<*xf32>

    %i8_min = "tf.Const"() {value = dense<-128.0> : tensor<f32>} : () -> tensor<f32>
    %i8_max = "tf.Const"() {value = dense<127.0> : tensor<f32>} : () -> tensor<f32>
    %clip = "tf.ClipByValue"(%round, %i8_min, %i8_max) : (tensor<*xf32>, tensor<f32>, tensor<f32>) -> tensor<*xf32>
    %i8 = "tf.Cast"(%clip) : (tensor<*xf32>) -> tensor<*xi8>
    func.return %i8 : tensor<*xi8>
  }

// CHECK: %[[QUANT_ARG:.*]] = stablehlo.uniform_quantize %[[ARG]] : (tensor<1x3x3x4xf32>) -> tensor<1x3x3x4x!quant.uniform<i8:f32, {{.*}}>>
// CHECK: %[[CONV:.*]] = stablehlo.convolution(%[[QUANT_ARG]], %[[FILTER]]) {{.*}} : (tensor<1x3x3x4x!quant.uniform<i8:f32, {{.*}}>>, tensor<3x3x4x4x!quant.uniform<i8:f32:3, {{.*}}>>) -> tensor<1x3x3x4x!quant.uniform<i8:f32, {{.*}}>>

    -> tensor<?x?x?x?xf32> {
  // CHECK-DAG: %[[EPS:.+]] = mhlo.constant dense<1.000000e-03> : tensor<f32>
  // CHECK-DAG: %[[VAR_SHAPE:.+]] = shape.shape_of %[[VARIANCE]] : tensor<?xf32> -> tensor<1xindex>
  // CHECK-DAG: %[[EPS_BCAST:.+]] =  "mhlo.dynamic_broadcast_in_dim"(%[[EPS]], %[[VAR_SHAPE]]) <{broadcast_dimensions = dense<> : tensor<0xi64>}> : (tensor<f32>, tensor<1xindex>) -> tensor<?xf32>

func.func @testIfCasts(%arg0: tensor<i1>, %arg1: tensor<!tf_type.variant<tensor<f32>>>) -> tensor<!tf_type.variant<tensor<f32>>> {
  %0 = "tf.If"(%arg0, %arg1) {
    then_branch = @testIfThen, else_branch = @testIfElse, is_stateless = false
  } : (tensor<i1>, tensor<!tf_type.variant<tensor<f32>>>) -> tensor<!tf_type.variant<tensor<f32>>>

    tf_device.return %12 : tensor<f32>
  }
  return %1#0, %1#1, %1#2, %1#3 : tensor<f32>, tensor<f32>, tensor<f32>, tensor<f32>
}
func.func private @_func(%arg0: tensor<f32>, %arg1: tensor<f32>, %arg2: tensor<f32>) -> tensor<f32> {
  %0 = "tf.Mul"(%arg0, %arg1) {device = "/device:TPU_REPLICATED_CORE:0"} : (tensor<f32>, tensor<f32>) -> tensor<f32>

  constexpr absl::string_view kFullyQuantizedAdd = R"mlir(
    func.func @fully_quantized_add(%arg0: tensor<2x!quant.uniform<i8:f32, 1.000000e+00:0>>) -> tensor<2x!quant.uniform<i8:f32, 1.000000e+00:0>> {
      %0 = stablehlo.add %arg0, %arg0 : tensor<2x!quant.uniform<i8:f32, 1.000000e+00:0>>
      return %0 : tensor<2x!quant.uniform<i8:f32, 1.000000e+00:0>>
    }
  )mlir";

  OwningOpRef<ModuleOp> module_op = ParseModuleOpString(kFullyQuantizedAdd);

  explicit TaggedValue(Func f32) : type_(FUNC), data_(f32) {}
  /// TaggedValue constructor for type FLOAT32.
  explicit TaggedValue(float f32) : type_(FLOAT32), data_(Float32(f32)) {}
  /// TaggedValue constructor for type INT64.
  explicit TaggedValue(int64_t i64) : type_(INT64), data_(Int64(i64)) {}
  /// TaggedValue constructor for type FLOAT32.
  explicit TaggedValue(Float32 f32) : type_(FLOAT32), data_(f32) {}