// CHECK:           return %[[VAL_5]] : tensor<8x8x8x64xf32>
// CHECK:         }
func.func @convert_conv2d_back_prop_input(%arg0: tensor<8x4x4x32xf32>, %arg1: tensor<3x3x64x32xf32>) -> tensor<8x8x8x64xf32> {
  %0 = "mhlo.convolution"(%arg0, %arg1) {batch_group_count = 1 : i64,
    dimension_numbers = #mhlo.conv<raw

  func.return %0#0, %1#0 : tensor<8x8x8x8xf32>, tensor<8x8x8x8xf32>
}
// CHECK: func @customAggregator
// CHECK-NEXT: %[[stats:.*]] = "quantfork.stats"(%arg0) <{layerStats = dense<[-1.000000e-01, 2.000000e-01]> : tensor<2xf32>}> : (tensor<8x8x8x8xf32>) -> tensor<8x8x8x8xf32>
// CHECK-NEXT: return %[[stats]], %arg0

    func.return %0#0 : tensor<8x8x8x8xf32>
  }

  // CHECK-LABEL: fusedBatchNormGradV3_Training

  func.return %0#0 : tensor<8x8x8x8xf32>
}

// -----

// CHECK-LABEL: fusedBatchNormGrad_Training

}

func.func @fakeQuantArgsTrue(%arg0: tensor<8x8x8x8xf32>) -> tensor<8x8x8x8xf32> {
  %0 = "tf.FakeQuantWithMinMaxArgs"(%arg0) {min = -0.1 : f32, max = 0.2 : f32, num_bits = 5, narrow_range = true} : (tensor<8x8x8x8xf32>) -> tensor<8x8x8x8xf32>
  func.return %0 : tensor<8x8x8x8xf32>

  // CHECK-LABEL: fakeQuantArgsTrue

  func.return %0 : tensor<256x16x32x32xf32>

  // LAYOUT-LABEL: Conv2dNCHW
  // LAYOUT: "tfl.conv_2d"
}

func.func @fusedBatchNormV3(tensor<8x8x8x8xf32>, tensor<8xf32>, tensor<8xf32>, tensor<8xf32>, tensor<8xf32>) -> (tensor<8x8x8x8xf32>, tensor<8xf32>) {
^bb0(%arg0: tensor<8x8x8x8xf32>, %arg1: tensor<8xf32>, %arg2: tensor<8xf32>, %arg3: tensor<8xf32>, %arg4: tensor<8xf32>):
  // OK

// permutated to match the new input shape.
// CHECK: (tensor<1x16x16x4xf32>, tensor<f32>) -> tensor<1x8x8x4xf32>

// Check that a `stablehlo.transpose` is added to the result to match the shape
// of the users.
// CHECK: %[[TRANSPOSE:.+]] = stablehlo.transpose %[[REDUCE_WINDOW]], dims = [0, 3, 1, 2] : (tensor<1x8x8x4xf32>) -> tensor<1x4x8x8xf32>
// CHECK: return %[[TRANSPOSE]]

// -----

// CHECK: return %[[v4]] : tensor<?x36x32xf32>
}

func.func @einsum_with_runtime_shape2(%arg0 : tensor<?x?x8x64xf32>, %arg1 : tensor<8x8x64xf32>) -> tensor<?x?x8xf32> {
  %0 = "tf.Einsum"(%arg0, %arg1) {device = "", equation = "ABNH,DNH->ABD"} : (tensor<?x?x8x64xf32>, tensor<8x8x64xf32>) -> tensor<?x?x8xf32>
  func.return %0 : tensor<?x?x8xf32>
// CHECK-LABEL: einsum_with_runtime_shape2

func.func @testFusedBatchNormToBatchNormV3(%arg0: tensor<8x8x8x8xf32>, %arg1: tensor<8xf32>, %arg2: tensor<8xf32>, %arg3: tensor<8xf32>, %arg4: tensor<8xf32>) -> (tensor<8x8x8x8xf32>) {
  // CHECK: "tf.FusedBatchNormV3"
  // CHECK-SAME: device = "/job:localhost/replica:0/task:0/device:GPU:0"