%3 = "tfl.quantize"(%2) {qtype = tensor<1x2x1x3x!quant.uniform<i8:f32, 1.0>>, volatile} : (tensor<1x2x1x3xf32>) -> tensor<1x2x1x3x!quant.uniform<i8:f32, 1.0>>
  %5 = "tfl.dequantize"(%arg0) : (tensor<1x2x1x3x!quant.uniform<i8:f64, 1.0>>) -> tensor<1x2x1x3xf32>
  "tfl.assign_variable"(%1, %5) : (tensor<!tf_type.resource>, tensor<1x2x1x3xf32>) -> ()

// RUN: tf-quant-opt %s -split-input-file -quant-insert-quantized-functions -quant-quantize-composite-functions | FileCheck %s

module {
  func.func @conv(%arg0: tensor<1x2x2x3xf32>) -> (tensor<*xf32>, tensor<*xf32>) {

    : (tensor<?x3x2x1xf32>, tensor<2x1x1x1xf32>) -> tensor<?x2x2x1xf32>
  %1 = chlo.broadcast_add %0, %zp_offset : (
      tensor<?x2x2x1xf32>, tensor<?x2x2x1xf32>) -> tensor<?x2x2x1xf32>
  %2 = chlo.broadcast_add %1, %bias : (
      tensor<?x2x2x1xf32>, tensor<1xf32>) ->tensor<?x2x2x1xf32>
  return %2 : tensor<?x2x2x1xf32>
}

// -----

// CHECK-LABEL: func @dot_general_add_add

func.func @depth_to_space(%arg0: tensor<1x1x1x4xf32>) -> tensor<1x2x2x1xf32> {
  %0 = "tf.DepthToSpace"(%arg0) {block_size = 2: i64,  data_format = "NHWC"}: (tensor<1x1x1x4xf32>) -> tensor<1x2x2x1xf32>
  func.return %0 : tensor<1x2x2x1xf32>
// CHECK: %[[CUSTOM_0:.*]] = "tfl.custom"(%arg0) <{custom_code = "FlexDepthToSpace", custom_option = #tfl<const_bytes : "{{.*}}">}> : (tensor<1x1x1x4xf32>) -> tensor<1x2x2x1xf32>
// CHECK: return %[[CUSTOM_0]] : tensor<1x2x2x1xf32>
}

  %3 = "tfl.quantize"(%2) {qtype = tensor<1x2x2x3x!quant.uniform<i8:f32, 0.1>>, volatile} : (tensor<1x2x2x3xf32>) -> tensor<1x2x2x3x!quant.uniform<i8:f32, 0.1>>
  %4 = "tfl.dequantize"(%3) : (tensor<1x2x2x3x!quant.uniform<i8:f32, 0.1>>) -> tensor<1x2x2x3xf32>

    %1 = "tf.BiasAdd"(%0, %arg2) {data_format = "NHWC", device = ""} : (tensor<?x2x2x2xf32>, tensor<2xf32>) -> tensor<?x2x2x2xf32>
    %2 = "tf.Relu6"(%1) {device = ""} : (tensor<?x2x2x2xf32>) -> tensor<?x2x2x2xf32>
    return %2 : tensor<?x2x2x2xf32>
  }

func.func @add_with_activation_transpose_permutation_mismatch(
      %arg0: tensor<1x2x3x4xf32>) -> tensor<1x3x2x4xf32> {
  %0 = stablehlo.constant dense<2.000000e+00> : tensor<1x3x2x4xf32>
  %1 = stablehlo.transpose %arg0, dims = [0, 2, 1, 3] : (tensor<1x2x3x4xf32>) -> tensor<1x3x2x4xf32>
  %2 = stablehlo.add %1, %0 : tensor<1x3x2x4xf32>
  return %2 : tensor<1x3x2x4xf32>
}
// CHECK: %[[TRANSPOSE_0:.+]] = stablehlo.transpose

    %6 = "quantfork.qcast"(%5) : (tensor<1x3x2x2xf32>) -> tensor<1x3x2x2x!quant.uniform<i8:f32, 0.0054901962186775953:-19>>
    %7 = "quantfork.dcast"(%6) : (tensor<1x3x2x2x!quant.uniform<i8:f32, 0.0054901962186775953:-19>>) -> tensor<1x3x2x2xf32>
    return %7 : tensor<1x3x2x2xf32>
  }

func.func @attribute_cast(%arg0: tensor<1x4x4x1xf32>) -> tensor<1x2x2x1xf32> {
  %0 = "tfr.cast"(%arg0) : (tensor<1x4x4x1xf32>) -> !tfr.tensor
  %stride_i32 = arith.constant 2 : i32
  %1 = tfr.call @tf__my_max_pool(%0, %stride_i32, %stride_i32) : (!tfr.tensor, i32, i32) -> !tfr.tensor
  %2 = "tfr.cast"(%1) : (!tfr.tensor) -> tensor<1x2x2x1xf32>
  func.return %2 : tensor<1x2x2x1xf32>
// CHECK: tf__max_pool
}