}
  return num_leading_broadcast_dims;
}

// Return the number of trailing 1s in the shape of the given input.
// Ex. input_shape = [1 x 1 x 2 x 1] => 1
// returns 0 if the input shape is not static.
int GetNumTrailingOnes(ShapedType input_type) {
  if (!input_type.hasStaticShape()) return 0;
  auto input_shape = input_type.getShape();
  int num_trailing_broadcast_dims = 0;
  for (int i = input_shape.size() - 1; i >= 0; --i) {