Repository: scipy/scipy
File: scipy/stats/_stats_py.py (Python)

Code before:
    loc_maxS = np.squeeze(np.take_along_axis(data_all, argmaxS, axis=-1), axis=-1)

    # Ensure sign of minS is not negative.
    minS = -np.squeeze(np.take_along_axis(cddiffs, argminS, axis=-1), axis=-1)
    maxS = np.squeeze(np.take_along_axis(cddiffs, argmaxS, axis=-1), axis=-1)
    minS = np.clip(minS, 0., 1.)

    if alternative == 'less':
        selector = np.ones(minS.shape, dtype=bool)
    elif alternative == 'two-sided':
        selector = minS > maxS
    else:
        selector = np.zeros(minS.shape, dtype=bool)

    d = np.where(selector, minS, maxS)
    d_location = np.where(selector, loc_minS, loc_maxS)
    d_sign = np.astype(np.where(selector, -1, 1), np.int8)
    prob = _ks_2samp_prob(d, n1, n2, mode, MAX_AUTO_N, alternative)

    if mode == 'exact' or (max(n1, n2) <= MAX_AUTO_N):
        g = np.gcd(n1, n2)
        lcm = (n1 // g) * n2
        h = np.round(d * lcm)
        d = h / lcm
        h_zero = (h == 0)
        prob[h_zero] = 1.0

    # Currently, `d` is a Python float. We want it to be a NumPy type, so
    # float64 is appropriate. An enhancement would be for `d` to respect the
    # dtype of the input.
    d = np.astype(d, np.float64, copy=False)
    return KstestResult(d[()], prob[()], statistic_location=d_location[()],
                        statistic_sign=d_sign[()])


@np.vectorize(excluded={1, 2, 3, 4, 5}, otypes=[np.float64])
def _ks_2samp_prob(d, n1, n2, mode, MAX_AUTO_N, alternative):
    g = np.gcd(n1, n2)
    n1g = n1 // g
    n2g = n2 // g
    prob = -np.inf
    if mode == 'auto':
        mode = 'exact' if max(n1, n2) <= MAX_AUTO_N else 'asymp'
    elif mode == 'exact':
        # If lcm(n1, n2) is too big, switch from exact to asymp
        if n1g >= np.iinfo(np.int32).max / n2g:
            mode = 'asymp'
            warnings.warn(
                f"Exact ks_2samp calculation not possible with samples sizes "
                f"{n1} and {n2}. Switching to 'asymp'.", RuntimeWarning,
                stacklevel=3)

Review comment (bug):
This code was (and still is) in `_attempt_exact_2kssamp`, and it modifies the test statistic. I moved it out of there because without it, a floating point equality check in `test_location_sign` started failing. I think it just cleans up the test statistic slightly, e.g. nudges from `np.float64(0.09999999999999998)` to `np.float64(0.1)`.

The funny thing is that this code was not running at all when the user passed `mode='asymp'`. In SciPy 1.16.2:
```python3
import numpy as np
from scipy import stats
x = np.arange(10, dtype=np.float64)
y = x.copy()
x[6] = 6.1
stats.ks_2samp(x, y, method='asymp').statistic  # np.float64(0.09999999999999998)
stats.ks_2samp(x, y, method='exact').statistic  # np.float64(0.1)
```
The choice of p-value calculation method shouldn't be changing the test statistic returned to the user, significantly or otherwise. So I suspect this code should either always run or never run. 

If it's OK with the reviewer, I'll just get rid of this here and change the equality check in the test to an `assert_allclose` with a tight tolerance. That will be helpful because `gcd` isn't in the array API.

```suggestion
```

Code after:

    # We want the ECDF of each sample evaluated at *all* the points in the pooled
    # sample. The values each ECDF can assume are given by:
    cdf1_vals = np.broadcast_to(np.linspace(0, 1, n1 + 1), batch_shape + (n1 + 1,))
    cdf2_vals = np.broadcast_to(np.linspace(0, 1, n2 + 1), batch_shape + (n2 + 1,))
    # Now we "just" need to know how many times each of these values *will* be assumed
    # when we evaluate the ECDFs at all points in the pooled sample.
    # These counts are given by the differences between consecutive ("min" or "max")
    # ranks corresponding with the observations in the (sorted) samples.
    ranks, data_all = _rankdata(data_all, method='min', return_sorted=True)  # axis=-1
    cdf1_counts = np.diff(ranks[..., :n1], prepend=1, append=n + 1, axis=-1)
    cdf2_counts = np.diff(ranks[..., -n2:], prepend=1, append=n + 1, axis=-1)
    # Repeat isn't vectorized - in general, this would produce a ragged array. However,
    # In our case, the sum of repeats for each slice is the same, so we can do a
    # vectorized repeat by raveling, repeating, then restoring the shape.
    cdf1 = np.repeat(np.ravel(cdf1_vals), np.ravel(cdf1_counts), axis=-1)
    cdf2 = np.repeat(np.ravel(cdf2_vals), np.ravel(cdf2_counts), axis=-1)
    cddiffs = np.reshape(cdf1 - cdf2, ranks.shape[:-1] + (-1,))

    # Identify the location of the statistic
    argminS = np.argmin(cddiffs, axis=-1, keepdims=True)
    argmaxS = np.argmax(cddiffs, axis=-1, keepdims=True)
    loc_minS = np.squeeze(np.take_along_axis(data_all, argminS, axis=-1), axis=-1)
    loc_maxS = np.squeeze(np.take_along_axis(data_all, argmaxS, axis=-1), axis=-1)

    # Ensure sign of minS is not negative.
    minS = -np.squeeze(np.take_along_axis(cddiffs, argminS, axis=-1), axis=-1)
    maxS = np.squeeze(np.take_along_axis(cddiffs, argmaxS, axis=-1), axis=-1)
    minS = np.clip(minS, 0., 1.)

    if alternative == 'less':
        selector = np.ones(minS.shape, dtype=bool)
    elif alternative == 'two-sided':
        selector = minS > maxS
    else:
        selector = np.zeros(minS.shape, dtype=bool)

    d = np.where(selector, minS, maxS)
    d_location = np.where(selector, loc_minS, loc_maxS)
    d_sign = np.where(selector, -1, 1).astype(np.int8)
    prob = _ks_2samp_prob(d, n1, n2, mode, MAX_AUTO_N, alternative)


    # Currently, `d` is a Python float. We want it to be a NumPy type, so
    # float64 is appropriate. An enhancement would be for `d` to respect the
    # dtype of the input.
    d = d.astype(np.float64, copy=False)
    return KstestResult(d[()], prob[()], statistic_location=d_location[()],
                        statistic_sign=d_sign[()])