return ranked && ranked.getRank() == (axis + 1) &&
             ranked.getShape()[axis] == 1;
    });
  };

  // Returns true if all binary ops operands at `operand_idx` index are scalars.
  auto is_all_scalars = [&](int operand_idx) -> bool {
    return llvm::all_of(op.getValues(), [&](Value arg) -> bool {
      if (exceptions.count(arg)) return true;
      auto operand = arg.getDefiningOp()->getOperand(operand_idx);