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

// CHECK: return %[[RES0]] : tensor<256x32x32x16xf32>

{
  "version": 3,
  "operator_codes": [
    {
      "builtin_code": "CONV_2D"
    }
  ],
  "subgraphs": [
    {

    %1 = "tf.PartitionedCall"(%arg0, %cst, %cst_0) {_tfl_quant_trait = "fully_quantizable", config = "", config_proto = "", executor_type = "", f = @composite_conv2d_with_bias_and_relu6_fn_1} : (tensor<1x2x2x3xf32>, tensor<2x2x3x2xf32>, tensor<2xf32>) -> tensor<*xf32> loc(callsit...

// CHECK: return %[[RES0]] : tensor<256x32x32x16xf32>

{
  "version": 3,
  "operator_codes": [
    {
      "builtin_code": "CONV_2D"
    }
  ],
  "subgraphs": [

    rhs_quantization_max_val = 127 : i64
  } : (tensor<1x2x2x3xf32>, tensor<2x3x3x2x!tf_type.qint8>, tensor<f32>, tensor<i32>) -> tensor<1x2x1x2xf32>
  func.return %0 : tensor<1x2x1x2xf32>
}

// -----

// CHECK-LABEL: func @uniform_quantized_convolution_hybrid_explicit
func.func @uniform_quantized_convolution_hybrid_explicit(%input: tensor<1x2x2x3xf32>) -> tensor<1x3x3x2xf32> {

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
}

  %res_p3 = "tfl.cumsum"(%arg, %axis_p3) {exclusive = false, reverse = false} : (tensor<1x2x1x3xf32>, tensor<i32>) -> tensor<1x2x1x3xf32>
  func.return %res_m4, %res_m3, %res_m2, %res_m1, %res_00, %res_p1, %res_p2, %res_p3 : tensor<1x2x1x3xf32>, tensor<1x2x1x3xf32>, tensor<1x2x1x3xf32>, tensor<1x2x1x3xf32>, tensor<1x2x1x3xf32>, tensor<1x2x1x3xf32>, tensor<1x2x1x3xf32>, tensor<1x2x1x3xf32>

// CHECK-NEXT:    return %[[VAL_2]] : tensor<1x2x4x1xf32>
// CHECK-NEXT:  }
}

// -----

module {
// expected-warning @+1 {{Invalid number of results from MaxUnpooling2D}}

// CHECK:         }

// -----

module {
func.func @constWeight(%arg0: tensor<256x32x32x3xf32>) -> tensor<256x30x30x16xf32> {
  %0 = "tfl.pseudo_const"() {value = dense<1.000000e+00> : tensor<16x3x3x3xf32>} : () -> tensor<16x3x3x3xf32>
  %1 = "tfl.pseudo_const"() {value = dense<1.000000e+00> : tensor<16xf32>} : () -> tensor<16xf32>

  %cst = arith.constant dense<[1, 3, 3, 2]> : tensor<4xi32>
  %cst_2 = arith.constant dense<[2, 3, 0, 1]> : tensor<4xi32>
  %2 = "tf.Transpose"(%arg, %cst_2) : (tensor<2x1x3x3xf32>, tensor<4xi32>) -> tensor<3x3x2x1xf32>
  %3 = "tf.Reshape"(%2, %cst) : (tensor<3x3x2x1xf32>, tensor<4xi32>) -> tensor<1x3x3x2xf32>
  return %3: tensor<1x3x3x2xf32>