Some reduction operators have two signatures, one that takes a dim and keepdim parameter to reduce along one or more dimensions and another that does not which reduces the flattened tensor. We should update these operators to only have one signature with an optional dim=None parameter which is more compatible with the Python Array API Standard.

In particular, the Python Array API Standard specifies the following types for axis (dim) parameter:
axis : Optional[ Union[ int, Tuple[ int, … ] ] ].

Operators with double signature to update

• argmax
• all
• any
• max
• min
• mean
• median
• nanmedian
• nansum
• prod
• std
• std_mean
• sum
• var
• var_mean
• count_nonzero [missing keepdim]

cc @mruberry @rgommers @pmeier @asmeurer @leofang @AnirudhDagar @asi1024 @emcastillo @kmaehashi @heitorschueroff