// PerChannel-NOT: tensor<2x3x3x2xf32>
// PerChannel-SAME: tensor<2x3x1x6xf32>
// PerChannel: %[[PARTITIONEDCALL_0:.*]] = "tf.PartitionedCall"(%arg0, %[[CONST_1:.*]]) <{config = "", config_proto = "", executor_type = "", f = @composite_depthwise_conv2d_fn_0}> {_tfl_quant_trait = "fully_quantizable"} : (tensor<1x3x4x3xf32>, tensor<2x3x1x6xf32>) -> tensor<*xf32>

// CHECK: %[[CONST_1:.*]] = arith.constant dense
// CHECK-NOT: tensor<2x3x3x2xf32>
// CHECK-SAME: tensor<2x3x1x6xf32>
// CHECK: %[[PARTITIONEDCALL_0:.*]] = "tf.PartitionedCall"(%arg0, %[[CONST_1:.*]]) <{config = "", config_proto = "", executor_type = "", f = @composite_depthwise_conv2d_fn_0}> {_tfl_quant_trait = "fully_quantizable"} : (tensor<1x3x4x3xf32>, tensor<2x3x1x6xf32>) -> tensor<*xf32>

    %dimension = "tf.Const"() { value = dense<3> : tensor<1xi64> } : () -> tensor<1xi64>
    %6 = "tf.Sum"(%3, %dimension) { keep_dims = true }: (tensor<1x3x2x2xf32>, tensor<1xi64>) -> tensor<1x3x2x1xf32>
    return %5, %6 : tensor<1x3x2x2xf32>, tensor<1x3x2x1xf32>
  }

// CHECK-LABEL: func @conv_with_multiple_uses
// CHECK: %[[div:.*]] = "tf.Div"(%arg0
// CHECK: %[[add:.*]] = "tf.AddV2"(%[[div]]
// CHECK: %[[maximum:.*]] = "tf.Maximum"(%[[add]]

// RUN: tf-mlir-translate -mlir-to-graphdef %s -tf-graph-as-function -o - | FileCheck %s

func.func @main(%arg0: tensor<*x!tf_type.resource>, %arg1: tensor<*x!tf_type.resource<tensor<3x3x1x32xf32>>>, %arg2: tensor<*xf32>, %arg3: tensor<2x4x6x8xi32>) -> (tensor<f32>, tensor<f32>)
attributes {tf.entry_function = {inputs = "args_0,args_1,args_2,args_3", outputs = "rets_0_RetVal,rets_1_RetVal"}} {
  %graph:2 = tf_executor.graph {

// Contains the `stablehlo.transpose` op of the arg (e.g. [b, f, 0, 1] to
// [b, 0, 1, f]). The weight constant is folded into [0, 1, i, o] format.
// CHECK-DAG: %[[CST:.+]] = stablehlo.constant dense<3.000000e+00> : tensor<3x3x8x8xf32>
// CHECK: %[[TRANSPOSE_1:.+]] = stablehlo.transpose %arg0, dims = [0, 2, 3, 1] : (tensor<1x8x4x4xf32>) -> tensor<1x4x4x8xf32>

// CHECK-LABEL: strips_attributes
// CHECK-NOT: tf
func.func @strips_attributes(%arg0: tensor<32x28x28x1xf32> {tf._user_specified_name = "x"},
                             %arg1: tensor<3x3x1x5xf32> {tf._user_specified_name = "w1"},
                             %arg2: tensor<5xf32> {tf._user_specified_name = "b1"},
                             %arg3: tensor<3920x10xf32> {tf._user_specified_name = "w2"},

// CHECK: %[[CONV:.+]] = stablehlo.convolution(%[[TRANSPOSE_0]], %[[TRANSPOSE_1]]) dim_numbers = [b, 0, 1, f]x[0, 1, i, o]->[b, 0, 1, f], window = {pad = {{\[\[}}1, 1], [1, 1]]} {batch_group_count = 1 : i64, feature_group_count = 1 : i64} : (tensor<1x4x4x8xf32>, tensor<3x3x8x8xf32>) -> tensor<1x4x4x8xf32>