// CHECK-NEXT: [[RESULT:%.*]] = "tf.BiasAdd"([[SQUEEZE]], [[BIAS]]) : (tensor<1x128x128xf32>, tensor<128xf32>) -> tensor<1x128x128xf32>
  // CHECK-NEXT: return [[RESULT]] : tensor<1x128x128xf32>
}

func.func @testDilatedDepthWiseConvWithExpandSqueeze1(%arg0: tensor<1x128x128xf32>, %arg1: tensor<5x5x1x1xf32>, %arg2: tensor<128xf32>) -> tensor<1x128x128xf32> {

  return %2 : tensor<1x1x1x4xf32>
}
func.func private @XlaCallModule_aten.avg_pool2d.default.impl_2(%arg0: tensor<1x1x1x8xf32>) -> tensor<1x1x1x4xf32>

// CHECK-LABEL: avg_pool2d_3
// CHECK: %cst = arith.constant dense<[0, 2, 3, 1]> : tensor<4xi32>
// CHECK: %0 = "tfl.transpose"(%arg0, %cst) : (tensor<1x1x1x8xf32>, tensor<4xi32>) -> tensor<1x1x8x1xf32>

  %0 = "tfl.fully_connected" (%arg0, %arg1, %arg2) {fused_activation_function = "NONE", keep_num_dims = false, weights_format = "DEFAULT"} : (tensor<1x256xf32>, tensor<128x256xf32>, tensor<128xf32>) -> tensor<1x128xf32>
  %1 = "tfl.relu"(%0) : (tensor<1x128xf32>) -> tensor<1x128xf32>
  func.return %1 : tensor<1x128xf32>

  // CHECK: %[[RES:[0-9].*]] = "tfl.fully_connected"

//   func.return %0#0 : tensor<1x1x2x8xf32>
// }
// COM: CHECK: func.func @main() -> tensor<1x1x2x8xf32>
// COM: CHECK-DAG: %[[CONST:.*]] = stablehlo.constant dense<{{.*}}> : tensor<1x1x2x8xf32>

func.func @fake_quant_with_min_max_vars(%arg0: tensor<1x1x28x48xf32>, %arg1: tensor<f32>, %arg2: tensor<f32>) -> tensor<1x1x28x48xf32> {
  %0 = "tf.FakeQuantWithMinMaxVars"(%arg0, %arg1, %arg2) {device = "", narrow_range = true, num_bits = 8 : i64} : (tensor<1x1x28x48xf32>, tensor<f32>, tensor<f32>) -> tensor<1x1x28x48xf32>
  func.return %0 : tensor<1x1x28x48xf32>
  // CHECK-SKIP: tf.FakeQuantWithMinMaxVars
  // CHECK-NOSKIP-NOT: tf.

// CHECK-NOT: "tfl.batch_matmul"
func.func @Batchmatmul2Fullyconnected(%arg0: tensor<4x128x2xf32>) -> (tensor<4x128x1xf32>) {
  %0 = arith.constant dense<[[1.0], [2.0]]> : tensor<2x1xf32>
  %1 = "tfl.batch_matmul"(%arg0, %0) {adj_x = false, adj_y = false, asymmetric_quantize_inputs = false} : (tensor<4x128x2xf32>, tensor<2x1xf32>) -> tensor<4x128x1xf32>
  func.return %1 : tensor<4x128x1xf32>
  // CHECK-NEXT: %[[CONST_WEIGHT:.*]] = arith.constant

       } : (tensor<1x32x32x3xf32>, tensor<4xi32>, tensor<1x32x32x8xf32>)
         -> tensor<1x1x3x8xf32>

  func.return %0 : tensor<1x1x3x8xf32>
}

// CHECK-LABEL: func @transposeConv2DBackpropInput
func.func @transposeConv2DBackpropInput(
  %input_sizes: tensor<4xi32>,
  %filter: tensor<1x1x3x8xf32>,
  %out_backprop: tensor<1x32x32x8xf32>
) -> tensor<1x32x32x3xf32> {

// -----

// CHECK-LABEL: testLiftDotConcatLHSAndRHS
func.func @testLiftDotConcatLHSAndRHS(%arg0: tensor<1x72x128xf32>, %arg1: tensor<1x128x72xf32>, %arg2: tensor<1x72x128xf32>, %arg3: tensor<1x128x72xf32>, %arg4: tensor<1x72x128xf32>, %arg5: tensor<1x128x72xf32>, %arg6: tensor<1x72x128xf32>, %arg7: tensor<1x128x72xf32>) -> tensor<4x72x72xf32> {
  %0 = "mhlo.dot_general"(%arg0, %arg1) {
    dot_dimension_numbers = #mhlo.dot<

// RUN: tf-opt %s -tf-layout-optimization=force-data-format=NCHW -verify-diagnostics | FileCheck %s --dump-input=always

// CHECK-LABEL: func @transposeConv2D
func.func @transposeConv2D(%arg0: tensor<1x3x32x32xf32>, %arg1: tensor<1x1x3x8xf32>) -> tensor<1x8x32x32xf32> {

  // Convert input: NCHW -> NHWC
  %0 = "tf.Const"() {value = dense<[0, 2, 3, 1]> : tensor<4xi32>} : () -> tensor<4xi32>

      stablehlo.return %3 : tensor<f32>
  }) {window_dimensions = array<i64: 1, 1, 2, 2>, window_strides = array<i64: 1, 1, 2, 2>} : (tensor<1x4x16x16xf32>, tensor<f32>) -> tensor<1x4x8x8xf32>
  return %2 : tensor<1x4x8x8xf32>
}
// CHECK-SAME: %[[ARG:.+]]: tensor<1x16x16x4xf32>
// CHECK-DAG: %[[INIT_VALUE_CONST:.+]] = stablehlo.constant dense<0xFF800000>

// Check that the body is not modified.