MLX compatibility: Statistical functions

[prady0t]

Array API	MLX Analog	Status	Notes	Test Node	Result
cumulative_sum(x, /, *, axis=None, dtype=None, include_initial=False)	cumsum(a, /, axis: int | None = None, *, reverse: bool = False, inclusive: bool = True, dtype: Dtype | None = None, stream=None)	incompatible	MLX lacks the include_initial keyword (raises TypeError) and skips the standard default-integer dtype promotion — a uint8 input stays uint8 instead of promoting to uint32.	test_cumulative_sum	Failed
cumulative_prod(x, /, *, axis=None, dtype=None, include_initial=False)	cumprod(a, /, axis: int | None = None, *, reverse: bool = False, inclusive: bool = True, dtype: Dtype | None = None, stream=None)	incompatible	Same two issues as cumulative_sum: missing include_initial keyword and no default-integer dtype promotion (uint8 stays uint8 instead of uint32).	test_cumulative_prod	Failed
prod(x, /, *, axis=None, dtype=None, keepdims=False)	prod(a, /, axis: None | int | Sequence[int] = None, keepdims: bool = False, *, stream=None)	incompatible	MLX has no dtype keyword (raises TypeError) and does not promote small integer inputs to the standard default integer dtype (uint8 stays uint8 instead of uint32).	test_prod	Failed
std(x, /, *, axis=None, correction=0.0, keepdims=False)	std(a, /, axis: None | int | Sequence[int] = None, keepdims: bool = False, ddof: int = 0, *, stream=None)	incompatible	MLX exposes ddof instead of the Array API's correction keyword, so passing correction raises a TypeError.	test_std	Failed
sum(x, /, *, axis=None, dtype=None, keepdims=False)	sum(a, /, axis: None | int | Sequence[int] = None, keepdims: bool = False, *, stream=None)	incompatible	MLX has no dtype keyword (raises TypeError) and does not promote small integer inputs to the standard default integer dtype (uint8 stays uint8 instead of uint32).	test_sum	Failed
var(x, /, *, axis=None, correction=0.0, keepdims=False)	var(a, /, axis: None | int | Sequence[int] = None, keepdims: bool = False, ddof: int = 0, *, stream=None)	incompatible	MLX exposes ddof instead of the Array API's correction keyword, so passing correction raises a TypeError.	test_var	Failed

[ev-br]

Thanks @prady0t.
Quick triage:

• the lack of type promotion (uint8 -> uint32 etc) looks systemic. Changing it has backwards compatibility implications, so I wouldn't be surprised if MLX rejects the change. Will be useful to discuss with them; I suggest to do it later though, after (i) small simple stuff is fixed, and (ii) once we know how other reductions behave. This is not something to change on a per-functions basis, I'd think.
• it's interesting to note that sum and prod accept tuple axis: this hints that moveaxis is a one-off; again, would be useful to see if other reductions accept tuple axes or not;
• ddof vs correction for var and std : is a small change, but has backwards compatibility implications. Worth asking upstream if they are open changing it.
• dtype keyword for sum and prod : is an enhancement without backwards compatibility concerns. Is a little tricky because the expectation is that the accumulation itself happens in the provided dtype. Still doable, so hopefully MLX accepts a PR.
• cumsum vs cumulative_sum: adding include_initial is a simple enhancement with no backcompat concerns. So is adding a new function, cumulative_sum an almost-alias for cumsum. What might be tricky is the dtype promotion story (keep consistent with other reductions or make consistent with the Array API speci). Maybe best to propose upstream after the dtype promotion, meaning after there's a yes or no decision on reductions in general.