module {
  func.func @matmul(%arg0: tensor<1x2x2x3xf32>) -> (tensor<*xf32>) {
    %cst_0 = "tf.Const"() {value = dense<0.000000e+00> : tensor<2x1024xf32>} : () -> tensor<2x1024xf32>
    %1 = "tf.PartitionedCall"(%arg0, %cst_0) {_tfl_quant_trait = "fully_quantizable", config = "", config_proto = "", executor_type = "", f = @composite_matmul_fn} : (tensor<1x2x2x3xf32>, tensor<2x1024xf32>) -> tensor<*xf32>
    func.return %1: tensor<*xf32>
  }

func.func private @conv(%input: tensor<1x3x4x3xf32> {tf._user_specified_name = "input_tensor"}) -> tensor<*xf32> attributes {tf._construction_context = "kEagerRuntime", tf._input_shapes = [#tf_type.shape<1x3x4x3>]} {
  %weight = arith.constant dense_resource<__elided__> : tensor<2x3x3x2xf32>
  %bias = arith.constant dense<[7.11401462, 7.05456924]> : tensor<2xf32>

// CHECK: %0 = tfl.sub %arg0, %cst_0 {fused_activation_function = "NONE"} : tensor<2x3x4x5xf32>
// CHECK: %1 = "tfl.transpose"(%0, %cst) : (tensor<2x3x4x5xf32>, tensor<4xi32>) -> tensor<5x2x3x4xf32>
// CHECK: return %1 : tensor<5x2x3x4xf32>

// -----

// CHECK-LABEL: permNotConstNoChange
func.func @permNotConstNoChange(%arg0: tensor<2x3x4x5xf32>, %perm: tensor<4xi32>) -> tensor<5x2x3x4xf32> {

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

// 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>) {

func.func @broadcast_to_to_reshape(%arg0: tensor<4x4x4xf32>, %arg1 : tensor<4xi32>) -> tensor<1x4x4x4xf32> {
  %0 = "tfl.broadcast_to"(%arg0, %arg1) : (tensor<4x4x4xf32>, tensor<4xi32>) -> tensor<1x4x4x4xf32>
  // CHECK: "tfl.reshape"
  // CHECK-SAME: (tensor<4x4x4xf32>, tensor<4xi32>) -> tensor<1x4x4x4xf32>
  func.return %0 : tensor<1x4x4x4xf32>
}

// Converts tfl.broadcast_to to tfl.reshape if input and output have the same

// -----

// CHECK-LABEL: transpose_simple_4d
func.func @transpose_simple_4d() -> tensor<5x2x3x4xf32> {
  %0 = stablehlo.constant dense<1.000000e+0> : tensor<2x3x4x5xf32>
  %1 = stablehlo.transpose %0, dims = [3, 0, 1, 2] : (tensor<2x3x4x5xf32>) -> tensor<5x2x3x4xf32>
  return %1 : tensor<5x2x3x4xf32>
}
// CHECK-DAG: %[[CONST_0:.+]] = stablehlo.constant dense<1.000000e+00> : tensor<5x2x3x4xf32>
// CHECK-NOT: transpose

// RUN: tf-quant-opt %s -split-input-file -quant-propagate-quantize-type | FileCheck %s

module {
  func.func @not_propagate_matmul(%arg0: tensor<1x2x2x2xf32>) -> tensor<*xf32> {
    %cst = "tf.Const"() {value = dense<127> : tensor<2x1024xi8>} : () -> tensor<2x1024xi8>
    %cst_0 = "tf.Const"() {value = dense<0.0157480314> : tensor<f32>} : () -> tensor<f32>
    %0 = "tf.Identity"(%cst) : (tensor<2x1024xi8>) -> tensor<2x1024xi8>

  %15 = "tf.Identity"(%14) {device = ""} : (tensor<1x2x4x1xf32>) -> tensor<1x2x4x1xf32>
  func.return %15 : tensor<1x2x4x1xf32>
}

// CHECK-LABEL: func @max_unpooling_2d(
// CHECK-SAME:                         %[[VAL_0:.*]]: tensor<1x1x2x1xf32>,
// CHECK-SAME:                         %[[VAL_1:.*]]: tensor<1x1x2x1xi32>) -> tensor<1x2x4x1xf32> attributes {tf._implements = "MaxUnpooling2D"} {

    %0 = "tf.Div"(%arg0, %cst_3) {device = ""} : (tensor<2x6x4x5xf32>, tensor<f32>) -> tensor<2x6x4x5xf32>
    %1 = "tf.AddV2"(%0, %cst_1) {device = ""} : (tensor<2x6x4x5xf32>, tensor<f32>) -> tensor<2x6x4x5xf32>
    %2 = "tf.Maximum"(%1, %cst_1) {device = ""} : (tensor<2x6x4x5xf32>, tensor<f32>) -> tensor<2x6x4x5xf32>
    %3 = "tf.Minimum"(%2, %cst_4) {device = ""} : (tensor<2x6x4x5xf32>, tensor<f32>) -> tensor<2x6x4x5xf32>