// CHECK{LITERAL}: %cst_1 = arith.constant dense<[[[[1.000000e+00], [2.000000e+00]], [[2.000000e+00], [4.000000e+00]]]]> : tensor<1x2x2x1xf32>
// CHECK: %3 = tfl.mul %2, %cst_1 {fused_activation_function = "NONE"} : tensor<1x2x2x1xf32>
// CHECK: %cst_2 = arith.constant dense<[0, 3, 1, 2]> : tensor<4xi32>

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
}

  %1 = "tf.Maximum"(%0, %cst_0) : (tensor<1x3x4x2xf32>, tensor<f32>) -> tensor<1x3x4x2xf32>
  %2 = "tf.Minimum"(%1, %cst_1) : (tensor<1x3x4x2xf32>, tensor<f32>) -> tensor<1x3x4x2xf32>
  func.return %2 : tensor<1x3x4x2xf32>

  %1 = "tf.Transpose"(%arg0, %0) : (tensor<1x4x4x8xf32>, tensor<4xi32>) -> tensor<1x8x4x4xf32>
  %2 = "tf.Transpose"(%1, %0) : (tensor<1x8x4x4xf32>, tensor<4xi32>) -> tensor<1x4x8x4xf32>
  %3 = "tf.Transpose"(%arg1, %0) : (tensor<1x4x4x8xf32>, tensor<4xi32>) -> tensor<1x8x4x4xf32>
  %4 = "tf.Transpose"(%3, %0) : (tensor<1x8x4x4xf32>, tensor<4xi32>) -> tensor<1x4x8x4xf32>
  %5 = "tf.AddV2"(%2, %4) : (tensor<1x4x8x4xf32>, tensor<1x4x8x4xf32>) -> tensor<1x4x8x4xf32>

  func.return %0 : tensor<2x4x4x1xf32>
}

func.func private @__inference_dense_image_warp(%arg0: tensor<2x4x4x1xf32>, %arg1: tensor<2x4x4x2xf32>) -> tensor<2x4x4x1xf32> attributes {tf._implements = "addons:DenseImageWarp"} {

// CHECK: %[[TRANSPOSE_2:.+]] = stablehlo.transpose %[[CUSTOM_AGGREGATOR_1]], dims = [0, 3, 1, 2] : (tensor<1x4x4x8xf32>) -> tensor<1x8x4x4xf32>

  return %2 : tensor<1x8x4x4xf32>
}

// CHECK-DAG: %[[CONST:.+]] = stablehlo.constant {{.*}} : tensor<8x8x3x3xf32>
// CHECK-DAG: %[[TRANSPOSE_0:.+]] = stablehlo.transpose %[[ARG]], dims = [0, 2, 3, 1] : (tensor<1x8x4x4xf32>) -> tensor<1x4x4x8xf32>

    return %1 : tensor<1x3x4x2xf32>
  }

  func.func private @composite_conv_fn(%arg0: tensor<1x3x4x3xf32>, %arg1: tensor<2x3x3x2xf32>) -> tensor<1x3x4x2xf32> attributes {_from_xla_call_module} {

  func.return %dequant : tensor<1x3x4x2xf32>
}

// -----

// CHECK-LABEL: @dont_merge_quantization_followed_by_quantization
// CHECK-SAME: %[[ARG_0:.*]]: tensor<1x3x4x3xf32>
func.func @dont_merge_quantization_followed_by_quantization(%arg0: tensor<1x3x4x3xf32>) -> tensor<1x3x4x3xf32> {
  // CHECK: %[[QUANT_ARG_0:.*]] = stablehlo.uniform_quantize %[[ARG_0]]

        _stablehlo_module_attrs = {}, _tfl_quant_trait = "fully_quantizable",
        device = ""
      } : (tensor<1x3x4x3xf32>, tensor<2x3x3x2xf32>) -> tensor<1x3x4x2xf32>
    return %0 : tensor<1x3x4x2xf32>
  }

  // CHECK: func.func private @qdq_for_conv_weight_per_channel_default(%[[ARG0:.+]]: tensor<1x3x4x3xf32>)
  // CHECK: %[[CST:.+]] = "tf.Const"() <{value = dense<3.000000e-01> : tensor<2x3x3x2xf32>}> : () -> tensor<2x3x3x2xf32>