ENH: stats: add array-API support to kstat/kstatvar

[j-bowhay]

towards #20544

[mdhaber]

Oops, I'm realizing that I probably shouldn't have put these on the initial list - they really should get native support for an axis argument before we add array API support. Would you add keyword argument axis with default None (for backward compatibility) to the signature and use it as appropriate?

[mdhaber]

Do you plan to add native axis support to these or should I do that?

[j-bowhay]

Do you plan to add native axis support to these or should I do that?

I was going to go for some of the lower-hanging fruit first so feel free to go for it in the mean time

[mdhaber]

Added in gh-20651. Sorry for the merge conflicts!

[mdhaber]

The conflicts are actually not as bad as I expected. It probably is worth it to merge and fix them rather than starting over. Please also convert the new tests added in gh-20651.

[j-bowhay]

Ok almost there just down to the final failures of this style:

scipy/stats/tests/test_morestats.py:1703: in test_nan_input
    xp_assert_equal(stats.kstat(data), xp.asarray(xp.nan))
        data       = Array([ 0.,  1.,  2.,  3.,  4.,  5., nan,
        7.,  8.,  9.], dtype=array_api_strict.float64)
        self       = <scipy.stats.tests.test_morestats.TestKstat object at 0x7f85c7677850>
        xp         = <module 'array_api_strict' from '/home/jakeb/miniconda3/envs/scipy-dev-pytorch/lib/python3.11/site-packages/array_api_strict/__init__.py'>
scipy/stats/_axis_nan_policy.py:405: in axis_nan_policy_wrapper
    return hypotest_fun_in(*args, **kwds)
        _no_deco   = False
        args       = (Array([ 0.,  1.,  2.,  3.,  4.,  5., nan,
        7.,  8.,  9.], dtype=array_api_strict.float64),)
        default_axis = None
        hypotest_fun_in = <function kstat at 0x7f8671774400>
        is_too_small = <function _axis_nan_policy_factory.<locals>.is_too_small at 0x7f86717742c0>
        kwd_samples = []
        kwds       = {}
        msg        = 'Use of `nan_policy` and `keepdims` is incompatible with non-NumPy arrays.'
        n_outputs  = 1
        n_samples  = 1
        override   = {'nan_propagation': True, 'vectorization': False}
        paired     = False
        result_to_tuple = <function <lambda> at 0x7f8671774220>
        temp       = Array([ 0.,  1.,  2.,  3.,  4.,  5., nan,
        7.,  8.,  9.], dtype=array_api_strict.float64)
        tuple_to_result = <function <lambda> at 0x7f8671774180>
scipy/stats/_morestats.py:307: in kstat
    S = [None] + [xp.sum(data**k, axis=axis) for k in range(1, n + 1)]
        N          = 10
        axis       = 0
        data       = array([ 0.,  1.,  2.,  3.,  4.,  5., nan,  7.,  8.,  9.])
        n          = 2
        xp         = <module 'array_api_strict' from '/home/jakeb/miniconda3/envs/scipy-dev-pytorch/lib/python3.11/site-packages/array_api_strict/__init__.py'>
scipy/stats/_morestats.py:307: in <listcomp>
    S = [None] + [xp.sum(data**k, axis=axis) for k in range(1, n + 1)]
        .0         = <range_iterator object at 0x7f85c75be580>
        axis       = 0
        data       = array([ 0.,  1.,  2.,  3.,  4.,  5., nan,  7.,  8.,  9.])
        k          = 1
        xp         = <module 'array_api_strict' from '/home/jakeb/miniconda3/envs/scipy-dev-pytorch/lib/python3.11/site-packages/array_api_strict/__init__.py'>
/home/jakeb/miniconda3/envs/scipy-dev-pytorch/lib/python3.11/site-packages/array_api_strict/_statistical_functions.py:100: in sum
    if x.dtype not in _numeric_dtypes:
        axis       = 0
        dtype      = None
        keepdims   = False
        x          = array([ 0.,  1.,  2.,  3.,  4.,  5., nan,  7.,  8.,  9.])
/home/jakeb/miniconda3/envs/scipy-dev-pytorch/lib/python3.11/site-packages/array_api_strict/_dtypes.py:24: in __eq__
    warnings.warn(
E   UserWarning: You are comparing a array_api_strict dtype against a NumPy native dtype object, but you probably don't want to do this. array_api_strict dtype objects compare unequal to their NumPy equivalents. Such cross-library comparison is not supported by the standard.
        other      = dtype('float64')
        self       = array_api_strict.float32
============================================================================================== short test summary info ==============================================================================================
FAILED scipy/stats/tests/test_morestats.py::TestKstat::test_moments_normal_distribution[array_api_strict] - UserWarning: You are comparing a array_api_strict dtype against a NumPy native dtype object, but you probably don't want to do this. array_api_strict dtype objects compare unequal to their NumPy equivalents. ...
FAILED scipy/stats/tests/test_morestats.py::TestKstat::test_nan_input[array_api_strict] - UserWarning: You are comparing a array_api_strict dtype against a NumPy native dtype object, but you probably don't want to do this. array_api_strict dtype objects compare unequal to their NumPy equivalents. ...

hopefully will be obvious when it isn't so late!

[mdhaber]

Line 296 and the next few don't look like they've been converted.

[j-bowhay]

Line 296 and the next few don't look like they've been converted.

Ha yes thanks, temporary post coursework submission blindness...

[mdhaber]

The statistic conversion looks pretty good, but the tests still need to become @array_api_compatible. We can add tests with the intent of checking the behavior for NumPy lists, but we also need to test behavior with backends other than NumPy.

[mdhaber]

After this, would you be willing to submit another PR that cleans up the notes of these? For example:

worse is that these correspond with:

but the implementation is:

[mdhaber]

Even torch will produce a float64 Tensor if it is generated from a NumPy float64 array, right?

[j-bowhay]

Yes, I can't quite remember why I did this, will have a look in the follow up

[mdhaber]

Fine for now, but nice to save @lucascolley some trouble by generating these without mutation as we go forward. Is there a better way than?

Suggested change

	data[6] = xp.nan
	data = xp.where(data==6, xp.nan, data)

[j-bowhay]

Obvously here this is fine but I do wonder if we encountered mutations in the in loop of a solver, for example, if all the new allocations might be a bit painful performance wise

[mdhaber]

Oh, I wouldn't change existing code like that all the time. Just easy cases like these tests, where it could make the difference between the test running with JAX or not. Cases do need to be considered individually.