) -> tensor<?x2x2x1xi32> {
  // CHECK-DAG: %[[conv:.*]] = mhlo.convolution
  // CHECK-DAG: %[[combined1:.*]] = chlo.broadcast_add %[[zp_offset:.*]], %[[bias:.*]]
  // CHECK-DAG: %[[combined2:.*]] = chlo.broadcast_add %[[combined1]], %[[bias]]
  // CHECK-DAG: %[[result:.*]] = chlo.broadcast_add %[[conv]], %[[combined2]]
  // CHECK: return %[[result]]
  %0 = mhlo.convolution(%lhs, %rhs)

  func.return %1: tensor<2xi32>
}

// CHECK-LABEL: func @broadcast_add
// TODO(laurenzo): Change this to a (5 + 2x1) shaped add to make the check
// patterns unambiguous and more interesting (once broadcastable trait is
// fixed upstream).
func.func @broadcast_add(%arg0: tensor<1xi32>, %arg1: tensor<1x2xi32>) -> tensor<1x2xi32> {

  // CHECK-DAG: %[[RHS:.*]] = mhlo.constant() <{value = dense<127> : tensor<2xi32>}> : () -> tensor<2x!quant.uniform<i32:f32, 2.000000e+00:4>>
  // CHECK: %[[RES:.*]] = chlo.broadcast_add %[[LHS]], %[[RHS]] {broadcast_dimensions = array<i64: 1>} :
  // CHECK-SAME: (tensor<3x2x!quant.uniform<i32:f32, 2.000000e+00:4>>, tensor<2x!quant.uniform<i32:f32, 2.000000e+00:4>>)