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

  return %0 : tensor<1x1x2x2xf32>
}
func.func private @XlaCallModule_aten.avg_pool2d.default.impl_4(%arg0: tensor<1x1x3x3xf32>) -> tensor<1x1x2x2xf32>

// CHECK-LABEL: avg_pool2d_5
// CHECK: %cst = arith.constant dense<[0, 2, 3, 1]> : tensor<4xi32>

  func.return %mm_s : tensor<1x3x3x2xf32>

// CHECK: %[[w:.*]] = arith.constant dense<1.270000e+02> : tensor<512x2xf32>
// CHECK: %[[q_w:.*]] = "tfl.quantize"(%[[w]]) <{qtype = tensor<512x2x!quant.uniform<i8<-127:127>:f32, 1.000000e+00>>}>
// CHECK: %[[dq_w:.*]] = "tfl.dequantize"(%[[q_w]]) : (tensor<512x2x!quant.uniform<i8<-127:127>:f32, 1.000000e+00>>) -> tensor<512x2xf32>

    %cst_1 = "tf.Const"() {device = "", value = dense<[-0.0291469581, 0.0106381178]> : tensor<2xf32>} : () -> tensor<2xf32>

  indices_are_sorted = false
} : (tensor<3x4x2xi32>, tensor<2x3x2xi64>) -> tensor<2x3x2x2xi32>
  func.return %2: tensor<2x3x2x2xi32>
}
// CHECK: %[[OPERAND:.*]] = stablehlo.constant
// CHECK: %[[INDICES:.*]] = stablehlo.constant
// CHECK: %[[XLA_CALL_MODULE:.*]] = "tf.XlaCallModule"(%[[OPERAND]], %[[INDICES]])
// CHECK: return %[[XLA_CALL_MODULE:.*]] : tensor<2x3x2x2xi32>
// CHECK: }

// CHECK-LABEL: private @composite_gather_fn_1

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

//===----------------------------------------------------------------------===//
// tf.UniformQuantize and tf.UniformDequantize legalization

}

// -----

// CHECK-LABEL: tensorlistConcat
func.func @tensorlistConcat(%arg0: tensor<3x2x2xf32>, %lead: tensor<i64>) -> (tensor<?x2xf32>, tensor<0xi64>) {
  %cst = arith.constant dense<[2, 2]> : tensor<2xi32>
  %list = "tf.TensorListFromTensor"(%arg0, %cst) : (tensor<3x2x2xf32>, tensor<2xi32>) -> tensor<!tf_type.variant<tensor<f32>>>

// %1 = stablehlo.constant dense<4>
// %2 = stablehlo.constant dense<2>
// %3 = stablehlo.convolution(%%arg0, %%arg1) :
//          (tensor<?x3x4x3xf32>, tensor<2x3x3x2xf32>) -> tensor<?x3x4x2xf32>
// %4 = stablehlo.get_dimension_size %3, dim = 0 :
//          (tensor<?x3x4x2xf32>) -> tensor<i32>
// %5 = stablehlo.reshape %4 :
//          (tensor<i32>) -> tensor<1xi32>
// %6 = stablehlo.concatenate %5, %0, %1, %2, dim = 0 :

    dilations = [1, 1, 1, 1],
    explicit_paddings = [],
    padding = "VALID",
    strides = [1, 1, 1, 1],
    use_cudnn_on_gpu = true
  } : (tensor<1x2x2x2xf32>, tensor<4xi32>, tensor<1x1x1x2xf32>) -> tensor<2x2x1x2xf32>
  func.return %result : tensor<2x2x1x2xf32>
}


// CHECK-LABEL: @conv3d_backprop_filter