ENH: stats.monte_carlo_test: add array API support

[mdhaber]

Reference issue

Towards gh-20544

What does this implement/fix?

This PR adds array-API support for scipy.stats.monte_carlo_test.

Additional information

Example (15~50x speedup locally, depending on dtype):

import numpy as np
import cupy as cp
from scipy import stats

rng = np.random.default_rng()
data_np = rng.standard_normal(size=100)
data_cp64 = cp.asarray(data_np, dtype=cp.float64)
data_cp32 = cp.asarray(data_np, dtype=cp.float32)

ref = stats.ttest_1samp(data_np, 0.)

def get_statistic(xp):
    def statistic(x, axis=0):
        m = xp.mean(x, axis=axis)
        # this is just the user-defined statistic;
        # it doesn't need to be array-API compliant
        v = xp.var(x, axis=axis, ddof=1)
        n = x.shape[axis]
        return m / (v / n) ** 0.5
    return statistic

# NumPy 64-bit
rng_np = np.random.default_rng()
res1 = stats.monte_carlo_test(data_np, rng_np.standard_normal, get_statistic(np), n_resamples=99999)
# 123 ms ± 1.02 ms per loop (mean ± std. dev. of 7 runs, 10 loops each)

# CuPy 64-bit
rng_cp = cp.random.default_rng()
res2 = stats.monte_carlo_test(data_cp64, rng_cp.standard_normal, get_statistic(cp), n_resamples=99999)
# 8.27 ms ± 767 µs per loop (mean ± std. dev. of 7 runs, 10 loops each)

# CuPy 32-bit
def rvs(size):
    return rng_cp.standard_normal(size=size, dtype=cp.float32)
res3 = stats.monte_carlo_test(data_cp32, rvs, get_statistic(cp), n_resamples=99999)
# 2.51 ms ± 33.4 µs per loop (mean ± std. dev. of 7 runs, 100 loops each)

print(ref.pvalue, res1.pvalue, res2.pvalue, res3.pvalue)
# 0.7506365948009475 0.754 0.75104 0.74768

[mdhaber]

Some self-review to aid the reviewer.

[mdhaber]

Again, didn't want to add to the mess of imports.

[rgommers]

Example (15~50x speedup locally, depending on dtype):

Very nice! This diff should be pretty clean once array-api-compat supports v2023.12 of the standard.

[mdhaber]

Thanks for taking a look @rgommers; responses are above. Another array API thing you might have thoughts about is the idea of making the default tolerance of xp_assert_close dtype-dependent.
#20597 (review)
The intent is to save us from constantly forcing torch arrays to be float64 to meet tolerances and/or setting tolerances for different dtypes manually.

[rgommers]

Towards #20544

If stats functions are going to be done one by one over many PRs rather than in one go, would it make sense to have a separate tracker for stats to point to? Having an overview of the coverage in stats will be useful (and help contributors pick functions to work on), and it's a bit much to put that overview in the scipy-wide tracker probably.

[mdhaber]

would it make sense to have a separate tracker for stats to point to?

? You referred to the separate stats tracker.

[rgommers]

? You referred to the separate stats tracker.

oops, sorry for the noise. temporary blindness there

[mdhaber]

@rgommers I know you'll be out for a bit - shall I ask others to review/merge this in the meantime, or did you want to continue with it?

[rgommers]

Thanks for asking. Please don't feel like you have to wait for me.

[mdhaber]

Replied to comments. Thanks @j-bowhay!

[lucascolley]

this all looks idiomatic to me, at a glance at least 👍

[lucascolley]

can we add a comment for why this is needed here too? (thanks for adding in the other places)

[j-bowhay]

Think this should be good to go once the merge conflicts and Lucas's comment is resolved

[mdhaber]

Thanks, all! I think I'll do power next, because it is quite similar and doesn't require an RNG.