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

  %0 = "mhlo.broadcast_in_dim"(%arg0) <{broadcast_dimensions = dense<[1, 2, 3]> : tensor<3xi64>, name = "broadcast.0"}> : (tensor<8x1x16xf32>) -> tensor<3x8x8x16xf32>
  func.return %0 : tensor<3x8x8x16xf32>
}

// CHECK-LABEL:   func @broadcast_in_dim_general_case(
// CHECK-SAME:                                        %[[VAL_0:.*]]: tensor<3x1x16xf32>) -> tensor<3x8x8x16xf32> {

  %cst_3 = "tf.Const"() {value = dense<[[[[1.400000e+01]], [[-2.800000e+01]], [[4.200000e+01]]], [[[-5.600000e+01]], [[7.100000e+01]], [[-8.500000e+01]]], [[[9.900000e+01]], [[-1.130000e+02]], [[1.270000e+02]]]]> : tensor<3x3x1x1xf32>} : () -> tensor<3x3x1x1xf32>
  %cst_4 = "tf.Const"() {value = dense<-1.280000e+02> : tensor<f32>} : () -> tensor<f32>
  %cst_5 = "tf.Const"() {value = dense<0.00118110236> : tensor<1xf32>} : () -> tensor<1xf32>

  %0 = "tf.Less"(%arg0, %arg1) : (tensor<8x7x6x5x?x3x2x1xf32>, tensor<?x3x2x1xf32>) -> tensor<8x7x6x5x?x3x2x1xi1>
  %1 = "tf.SelectV2"(%0, %arg0, %arg1) : (tensor<8x7x6x5x?x3x2x1xi1>, tensor<8x7x6x5x?x3x2x1xf32>, tensor<?x3x2x1xf32>) -> tensor<8x7x6x5x?x3x2x1xf32>
  func.return %1 : tensor<8x7x6x5x?x3x2x1xf32>

  return %2 : tensor<?x2x2x1xi8>
}

// -----

// CHECK-LABEL: func @convolution_add_add_f32
func.func @convolution_add_add_f32(
    %lhs: tensor<?x3x2x1xf32>, %rhs: tensor<2x1x1x1xf32>,
    %zp_offset: tensor<?x2x2x1xf32>, %bias: tensor<1xf32>
  ) -> tensor<?x2x2x1xf32> {
  // CHECK-DAG: %[[conv:.*]] = mhlo.convolution

// CHECK: Number of dequantize layers added: 1
}


// -----

module {
  func.func @float_einsum(%arg0: tensor<?x64x32xf32>, %arg1: tensor<32x2x16xf32>) -> (tensor<?x64x2x16xf32>) {
    %0 = "tf.Einsum"(%arg0, %arg1) {equation = "abc,cde->abde"} : (tensor<?x64x32xf32>, tensor<32x2x16xf32>) -> tensor<?x64x2x16xf32>
    func.return %0 : tensor<?x64x2x16xf32>
  }

// CHECK-LABEL: func @float_einsum

  func.func @conv2d_partially_dynamic_spatial_dim(%arg0: tensor<256x?x32x3xf32>, %arg1: tensor<3x3x3x16xf32>) -> tensor<*xf32> {
    // CHECK: "tf.Conv2D"
    // CHECK-SAME: -> tensor<256x?x32x16xf32>
    %0 = "tf.Conv2D"(%arg0, %arg1) {padding = "SAME", strides = [1, 1, 1, 1]} : (tensor<256x?x32x3xf32>, tensor<3x3x3x16xf32>) -> tensor<*xf32>
    func.return %0 : tensor<*xf32>
  }

  %0 = "tf.Const"() {value = dense<0.111111112> : tensor<3x3x1x1xf32>} : () -> tensor<3x3x1x1xf32>
  %1 = "tf.Const"() {value = dense<1.000000e+00> : tensor<1x520x520x1xf32>} : () -> tensor<1x520x520x1xf32>
  %2 = "tf.DepthwiseConv2dNative"(%1, %0) {data_format = "NHWC", device = "", dilations = [1, 1, 1, 1], explicit_paddings = [], padding = "SAME", strides = [1, 1, 1, 1]} : (tensor<1x520x520x1xf32>, tensor<3x3x1x1xf32>) -> tensor<1x520x520x1xf32>

  // CHECK: mhlo.convolution(%arg0, %arg1)
  // CHECK-SAME{LITERAL}: pad = [[6, 0], [3, 3]]

  %0 = "tf.Conv2D"(%arg0, %arg1) {data_format = "NHWC", dilations = [1, 2, 3, 1], padding = "EXPLICIT", explicit_paddings = [0, 0, 6, 0, 3, 3, 0, 0], strides = [1, 4, 5, 1]} : (tensor<256x32x32x6xf32>, tensor<3x3x3x16xf32>) -> tensor<256x9x7x16xf32>

  func.return %1 : tensor<256x30x30x16xf32>

// CHECK-DAG: %[[w:.*]] = arith.constant dense<1.000000e+00> : tensor<3x3x3x16xf32>