Repository: scipy/scipy
Pull request: ENH: stats.pearsonr: add array API support
#### Reference issue
gh-20137
Closes gh-20324

#### What does this implement/fix?
This explores the addition of Array API support to `scipy.stats.pearsonr`. Only the last commit is relevant; the others are from gh-20137.

#### Additional information
Need to resolve merge conflicts and skip a test on 32-bit, but otherwise, I think this is ready to go.

---
*Old news:*

Most of this will be pretty straightforward. There are some little things I'll want to address later (e.g. previously, `pearsonr` converted inputs to be at least `float64`, but I imagine we'd want to respect dtype with array API), but there is one big question for now:

Calculation of the p-value currently relies on the incomplete beta function, which is not among the special functions for which we have experimental array API support (gh-19023). Even if it were, calculation of the p-value currently relies on the distribution infrastructure. In any case, calculation of the statistic with an alternative array backend can easily be done now, but calculation of the p-value with an alternative array backend will take more time. 

For vectorized calculations, I think there is still value in calculating the statistic with the alternative array backend, then converting the statistic (which has been reduced by at least one dimension) to a NumPy array for calculation of the p-value. Are there objections to this?

I know there have been objections to converting non-NumPy arrays to NumPy arrays for running compiled code, but I think this is a little different, since the statistic might not even be an array after the reducing operation, and if it is, it's smaller than the original array.

Problem statement:
MAINT, BUG (?): pearsonr statistic return type change
Originally reported by @stefanv on mailing list, for an issue that was detected in `skimage`. My local testing suggests that the issue is reproducible for NumPy `1.26.4` and NumPy `2.0.0b1` though, so the reason may be due to a code change we've made directly.

I confirmed locally that the test below passes at `v1.12.0` tag, and fails on `main`. I think Stefan's argument was that the docstring still says `float` for the statistic.

```diff
--- a/scipy/stats/tests/test_stats.py
+++ b/scipy/stats/tests/test_stats.py
@@ -8495,3 +8495,8 @@ class TestExpectile:
         for alpha in np.r_[0, alpha_seq, 1 - alpha_seq[:-1:-1], 1]:
             e_list.append(stats.expectile(x, alpha=alpha))
         assert np.all(np.diff(e_list) > 0)
+
+
+def test_pearsonr_stat_ret_type():
+    res = stats.pearsonr([1, 2, 3], [4, 5, 6])
+    assert isinstance(res.statistic, float)
```

Maybe one of the `stats` folks can comment re: what we might want to do here for `1.13.0`? I'll see if I can bisect this for you folks.

Hints from issue discussion:

Commit message: stats.pearsonr: add Array API support

Patch to review:
diff --git a/scipy/stats/_morestats.py b/scipy/stats/_morestats.py
index 8da934504184..13d6824f7461 100644
--- a/scipy/stats/_morestats.py
+++ b/scipy/stats/_morestats.py
@@ -1250,6 +1250,12 @@ def boxcox_normmax(
     6.000...
     """
     x = np.asarray(x)
+
+    if np.any(~np.isfinite(x) | (x < 0)):
+        message = ("The `x` argument of `boxcox_normmax` must contain "
+                   "only positive, finite, real numbers.")
+        raise ValueError(message)
+
     end_msg = "exceed specified `ymax`."
     if isinstance(ymax, _BigFloat):
         dtype = x.dtype if np.issubdtype(x.dtype, np.floating) else np.float64
@@ -1321,15 +1327,7 @@ def _all(x):
 
     optimfunc = methods[method]
 
-    try:
-        res = optimfunc(x)
-    except ValueError as e:
-        if "infs or NaNs" in str(e):
-            message = ("The `x` argument of `boxcox_normmax` must contain "
-                       "only positive, finite, real numbers.")
-            raise ValueError(message) from e
-        else:
-            raise e
+    res = optimfunc(x)
 
     if res is None:
         message = ("The `optimizer` argument of `boxcox_normmax` must return "
diff --git a/scipy/stats/_stats_py.py b/scipy/stats/_stats_py.py
index a0bcb13c74b2..836daefeff06 100644
--- a/scipy/stats/_stats_py.py
+++ b/scipy/stats/_stats_py.py
@@ -45,7 +45,9 @@
                               AxisError)
 
 import scipy.special as special
-from scipy import linalg
+# Import unused here but needs to stay until end of deprecation periode
+# See https://github.com/scipy/scipy/issues/15765#issuecomment-1875564522
+from scipy import linalg  # noqa: F401
 from . import distributions
 from . import _mstats_basic as mstats_basic
 from ._stats_mstats_common import (_find_repeats, linregress, theilslopes,
@@ -66,6 +68,7 @@
 from scipy.optimize import root_scalar
 from scipy._lib.deprecation import _NoValue, _deprecate_positional_args
 from scipy._lib._util import normalize_axis_index
+from scipy._lib._array_api import array_namespace, is_numpy
 
 # In __all__ but deprecated for removal in SciPy 1.13.0
 from scipy._lib._util import float_factorial  # noqa: F401
@@ -4347,47 +4350,47 @@ def _pearsonr_fisher_ci(r, n, confidence_level, alternative):
     Fisher's transformation is used to compute the confidence interval
     (https://en.wikipedia.org/wiki/Fisher_transformation).
     """
-    if r == 1:
-        zr = np.inf
-    elif r == -1:
-        zr = -np.inf
-    else:
-        zr = np.arctanh(r)
+    xp = array_namespace(r)
+
+    with np.errstate(divide='ignore'):
+        zr = xp.arctanh(r)
 
+    ones = xp.ones_like(r)
+    n, confidence_level = xp.asarray([n, confidence_level], dtype=r.dtype)
     if n > 3:
-        se = np.sqrt(1 / (n - 3))
+        se = xp.sqrt(1 / (n - 3))
         if alternative == "two-sided":
             h = special.ndtri(0.5 + confidence_level/2)
             zlo = zr - h*se
             zhi = zr + h*se
-            rlo = np.tanh(zlo)
-            rhi = np.tanh(zhi)
+            rlo = xp.tanh(zlo)
+            rhi = xp.tanh(zhi)
         elif alternative == "less":
             h = special.ndtri(confidence_level)
             zhi = zr + h*se
-            rhi = np.tanh(zhi)
-            rlo = -1.0
+            rhi = xp.tanh(zhi)
+            rlo = -ones
         else:
             # alternative == "greater":
             h = special.ndtri(confidence_level)
             zlo = zr - h*se
-            rlo = np.tanh(zlo)
-            rhi = 1.0
+            rlo = xp.tanh(zlo)
+            rhi = ones
     else:
-        rlo, rhi = -1.0, 1.0
+        rlo, rhi = -ones, ones
 
-    return ConfidenceInterval(low=rlo, high=rhi)
+    return ConfidenceInterval(low=rlo[()], high=rhi[()])
 
 
-def _pearsonr_bootstrap_ci(confidence_level, method, x, y, alternative):
+def _pearsonr_bootstrap_ci(confidence_level, method, x, y, alternative, axis):
     """
     Compute the confidence interval for Pearson's R using the bootstrap.
     """
-    def statistic(x, y):
-        statistic, _ = pearsonr(x, y)
+    def statistic(x, y, axis):
+        statistic, _ = pearsonr(x, y, axis=axis)
         return statistic
 
-    res = bootstrap((x, y), statistic, confidence_level=confidence_level,
+    res = bootstrap((x, y), statistic, confidence_level=confidence_level, axis=axis,
                     paired=True, alternative=alternative, **method._asdict())
     # for one-sided confidence intervals, bootstrap gives +/- inf on one side
     res.confidence_interval = np.clip(res.confidence_interval, -1, 1)
@@ -4419,12 +4422,13 @@ class PearsonRResult(PearsonRResultBase):
         coefficient `statistic` for the given confidence level.
 
     """
-    def __init__(self, statistic, pvalue, alternative, n, x, y):
+    def __init__(self, statistic, pvalue, alternative, n, x, y, axis):
         super().__init__(statistic, pvalue)
         self._alternative = alternative
         self._n = n
         self._x = x
         self._y = y
+        self._axis = axis
 
         # add alias for consistency with other correlation functions
         self.correlation = statistic
@@ -4475,8 +4479,14 @@ def confidence_interval(self, confidence_level=0.95, method=None):
                https://en.wikipedia.org/wiki/Pearson_correlation_coefficient
         """
         if isinstance(method, BootstrapMethod):
-            ci = _pearsonr_bootstrap_ci(confidence_level, method,
-                                        self._x, self._y, self._alternative)
+            xp = array_namespace(self._x)
+            message = ('`method` must be None if `pearsonr` '
+                       'arguments were not NumPy arrays.')
+            if not is_numpy(xp):
+                raise ValueError(message)
+
+            ci = _pearsonr_bootstrap_ci(confidence_level, method, self._x, self._y,
+                                        self._alternative, self._axis)
         elif method is None:
             ci = _pearsonr_fisher_ci(self.statistic, self._n, confidence_level,
                                      self._alternative)
@@ -4486,7 +4496,7 @@ def confidence_interval(self, confidence_level=0.95, method=None):
             raise ValueError(message)
         return ci
 
-def pearsonr(x, y, *, alternative='two-sided', method=None):
+def pearsonr(x, y, *, alternative='two-sided', method=None, axis=0):
     r"""
     Pearson correlation coefficient and p-value for testing non-correlation.
 
@@ -4508,10 +4518,15 @@ def pearsonr(x, y, *, alternative='two-sided', method=None):
 
     Parameters
     ----------
-    x : (N,) array_like
+    x : array_like
         Input array.
-    y : (N,) array_like
+    y : array_like
         Input array.
+    axis : int or None, default
+        Axis along which to perform the calculation. Default is 0.
+        If None, ravel both arrays before performing the calculation.
+
+        .. versionadded:: 1.13.0
     alternative : {'two-sided', 'greater', 'less'}, optional
         Defines the alternative hypothesis. Default is 'two-sided'.
         The following options are available:
@@ -4669,6 +4684,24 @@ def pearsonr(x, y, *, alternative='two-sided', method=None):
     >>> res.confidence_interval(confidence_level=0.9, method=method)
     ConfidenceInterval(low=-0.9983163756488651, high=-0.22771001702132443)  # may vary
 
+    If N-dimensional arrays are provided, multiple tests are performed in a
+    single call according to the same conventions as most `scipy.stats` functions:
+
+    >>> rng = np.random.default_rng(2348246935601934321)
+    >>> x = rng.standard_normal((8, 15))
+    >>> y = rng.standard_normal((8, 15))
+    >>> stats.pearsonr(x, y, axis=0).statistic.shape  # between corresponding columns
+    (15,)
+    >>> stats.pearsonr(x, y, axis=1).statistic.shape  # between corresponding rows
+    (8,)
+
+    To perform all pairwise comparisons between slices of the arrays,
+    use standard NumPy broadcasting techniques. For instance, to compute the
+    correlation between all pairs of rows:
+
+    >>> stats.pearsonr(x[:, np.newaxis, :], y, axis=-1).statistic.shape
+    (8, 8)
+
     There is a linear dependence between x and y if y = a + b*x + e, where
     a,b are constants and e is a random error term, assumed to be independent
     of x. For simplicity, assume that x is standard normal, a=0, b=1 and let
@@ -4716,85 +4749,118 @@ def pearsonr(x, y, *, alternative='two-sided', method=None):
     than zero which happens in about half of the cases if we sample x and y.
 
     """
-    n = len(x)
-    if n != len(y):
-        raise ValueError('x and y must have the same length.')
+    xp = array_namespace(x, y)
+    x = xp.asarray(x)
+    y = xp.asarray(y)
+
+    if not is_numpy(xp) and method is not None:
+        method = 'invalid'
+
+    if axis is None:
+        x = xp.reshape(x, (-1,))
+        y = xp.reshape(y, (-1,))
+        axis = -1
+
+    axis_int = int(axis)
+    if axis_int != axis:
+        raise ValueError('`axis` must be an integer.')
+    axis = axis_int
+
+    n = x.shape[axis]
+    if n != y.shape[axis]:
+        raise ValueError('`x` and `y` must have the same length along `axis`.')
 
     if n < 2:
-        raise ValueError('x and y must have length at least 2.')
+        raise ValueError('`x` and `y` must have length at least 2.')
 
-    x = np.asarray(x)
-    y = np.asarray(y)
+    try:
+        x, y = xp.broadcast_arrays(x, y)
+    except ValueError as e:
+        message = '`x` and `y` must be broadcastable.'
+        raise ValueError(message) from e
+
+    x = xp.moveaxis(x, axis, -1)
+    y = xp.moveaxis(y, axis, -1)
+    axis = -1
+
+    dtype = xp.result_type(x.dtype, y.dtype)
+    if xp.isdtype(dtype, "integral"):
+        dtype = xp.float64
 
-    if (np.issubdtype(x.dtype, np.complexfloating)
-            or np.issubdtype(y.dtype, np.complexfloating)):
+    if xp.isdtype(dtype, "complex floating"):
         raise ValueError('This function does not support complex data')
 
+    x = xp.astype(x, dtype, copy=False)
+    y = xp.astype(y, dtype, copy=False)
+    threshold = xp.finfo(dtype).eps ** 0.75
+
     # If an input is constant, the correlation coefficient is not defined.
-    if (x == x[0]).all() or (y == y[0]).all():
+    const_x = xp.all(x == x[..., [0]], axis=-1)
+    const_y = xp.all(y == y[..., [0]], axis=-1)
+    const_xy = const_x | const_y
+    if xp.any(const_xy):
         msg = ("An input array is constant; the correlation coefficient "
                "is not defined.")
         warnings.warn(stats.ConstantInputWarning(msg), stacklevel=2)
-        result = PearsonRResult(statistic=np.nan, pvalue=np.nan, n=n,
-                                alternative=alternative, x=x, y=y)
-        return result
 
     if isinstance(method, PermutationMethod):
-        def statistic(y):
-            statistic, _ = pearsonr(x, y, alternative=alternative)
+        def statistic(y, axis):
+            statistic, _ = pearsonr(x, y, axis=axis, alternative=alternative)
             return statistic
 
         res = permutation_test((y,), statistic, permutation_type='pairings',
-                               alternative=alternative, **method._asdict())
+                               axis=axis, alternative=alternative, **method._asdict())
 
         return PearsonRResult(statistic=res.statistic, pvalue=res.pvalue, n=n,
-                              alternative=alternative, x=x, y=y)
+                              alternative=alternative, x=x, y=y, axis=axis)
     elif isinstance(method, MonteCarloMethod):
-        def statistic(x, y):
-            statistic, _ = pearsonr(x, y, alternative=alternative)
+        def statistic(x, y, axis):
+            statistic, _ = pearsonr(x, y, axis=axis, alternative=alternative)
             return statistic
 
         if method.rvs is None:
             rng = np.random.default_rng()
             method.rvs = rng.normal, rng.normal
 
-        res = monte_carlo_test((x, y,), statistic=statistic,
+        res = monte_carlo_test((x, y,), statistic=statistic, axis=axis,
                                alternative=alternative, **method._asdict())
 
         return PearsonRResult(statistic=res.statistic, pvalue=res.pvalue, n=n,
-                              alternative=alternative, x=x, y=y)
+                              alternative=alternative, x=x, y=y, axis=axis)
+    elif method == 'invalid':
+        message = ('`method` must be None if arguments are not NumPy arrays.')
+        raise ValueError(message)
     elif method is not None:
         message = ('`method` must be an instance of `PermutationMethod`,'
                    '`MonteCarloMethod`, or None.')
         raise ValueError(message)
 
-    # dtype is the data type for the calculations.  This expression ensures
-    # that the data type is at least 64 bit floating point.  It might have
-    # more precision if the input is, for example, np.longdouble.
-    dtype = type(1.0 + x[0] + y[0])
-
     if n == 2:
-        r = dtype(np.sign(x[1] - x[0])*np.sign(y[1] - y[0]))
-        result = PearsonRResult(statistic=r, pvalue=1.0, n=n,
-                                alternative=alternative, x=x, y=y)
+        r = (xp.sign(x[..., 1] - x[..., 0])*xp.sign(y[..., 1] - y[..., 0]))[()]
+        result = PearsonRResult(statistic=r, pvalue=xp.ones_like(r)[()], n=n,
+                                alternative=alternative, x=x, y=y, axis=axis)
         return result
 
-    xmean = x.mean(dtype=dtype)
-    ymean = y.mean(dtype=dtype)
-
-    # By using `astype(dtype)`, we ensure that the intermediate calculations
-    # use at least 64 bit floating point.
-    xm = x.astype(dtype) - xmean
-    ym = y.astype(dtype) - ymean
+    xmean = x.mean(axis=axis, keepdims=True)
+    ymean = y.mean(axis=axis, keepdims=True)
+    xm = x - xmean
+    ym = y - ymean
 
-    # Unlike np.linalg.norm or the expression sqrt((xm*xm).sum()),
-    # scipy.linalg.norm(xm) does not overflow if xm is, for example,
+    # scipy.linalg.norm(xm) avoids premature overflow when xm is e.g.
     # [-5e210, 5e210, 3e200, -3e200]
-    normxm = linalg.norm(xm)
-    normym = linalg.norm(ym)
-
-    threshold = 1e-13
-    if normxm < threshold*abs(xmean) or normym < threshold*abs(ymean):
+    # but not when `axis` is provided, so scale manually. scipy.linalg.norm
+    # also raises an error with NaN input rather than returning NaN, so
+    # use np.linalg.norm.
+    xmax = xp.max(xp.abs(xm), axis=axis, keepdims=True)
+    ymax = xp.max(xp.abs(ym), axis=axis, keepdims=True)
+    with np.errstate(invalid='ignore'):
+        normxm = xmax * xp.linalg.vector_norm(xm/xmax, axis=axis, keepdims=True)
+        normym = ymax * xp.linalg.vector_norm(ym/ymax, axis=axis, keepdims=True)
+
+    nconst_x = xp.any(normxm < threshold*xp.abs(xmean), axis=axis)
+    nconst_y = xp.any(normym < threshold*xp.abs(ymean), axis=axis)
+    nconst_xy = nconst_x | nconst_y
+    if xp.any(nconst_xy & (~const_xy)):
         # If all the values in x (likewise y) are very close to the mean,
         # the loss of precision that occurs in the subtraction xm = x - xmean
         # might result in large errors in r.
@@ -4802,20 +4868,25 @@ def statistic(x, y):
                "correlation coefficient may be inaccurate.")
         warnings.warn(stats.NearConstantInputWarning(msg), stacklevel=2)
 
-    r = np.dot(xm/normxm, ym/normym)
+    with np.errstate(invalid='ignore', divide='ignore'):
+        r = xp.sum(xm/normxm * ym/normym, axis=axis)
 
     # Presumably, if abs(r) > 1, then it is only some small artifact of
     # floating point arithmetic.
-    r = max(min(r, 1.0), -1.0)
+    one = xp.asarray(1, dtype=dtype)
+    r = xp.asarray(xp.clip(r, -one, one))
+    r[const_xy] = xp.asarray(np.nan, dtype=dtype)
 
     # As explained in the docstring, the distribution of `r` under the null
     # hypothesis is the beta distribution on (-1, 1) with a = b = n/2 - 1.
+    # This needs to be done with NumPy arrays given the existing infrastructure.
     ab = n/2 - 1
     dist = stats.beta(ab, ab, loc=-1, scale=2)
-    pvalue = _get_pvalue(r, dist, alternative)
+    pvalue = _get_pvalue(np.asarray(r), dist, alternative)
+    pvalue = xp.asarray(pvalue, dtype=dtype)
 
-    return PearsonRResult(statistic=r, pvalue=pvalue, n=n,
-                          alternative=alternative, x=x, y=y)
+    return PearsonRResult(statistic=r[()], pvalue=pvalue[()], n=n,
+                          alternative=alternative, x=x, y=y, axis=axis)
 
 
 def fisher_exact(table, alternative='two-sided'):
diff --git a/scipy/stats/tests/test_stats.py b/scipy/stats/tests/test_stats.py
index a1600a95e560..20ce0fc6f7ce 100644
--- a/scipy/stats/tests/test_stats.py
+++ b/scipy/stats/tests/test_stats.py
@@ -38,6 +38,8 @@
 from scipy.stats._axis_nan_policy import _broadcast_concatenate
 from scipy.stats._stats_py import _permutation_distribution_t
 from scipy._lib._util import AxisError
+from scipy._lib._array_api import xp_assert_close
+from scipy.conftest import array_api_compatible
 
 
 """ Numbers in docstrings beginning with 'W' refer to the section numbers
@@ -401,7 +403,7 @@ def test_very_small_input_values(self):
         # See https://github.com/scipy/scipy/issues/9353
         x = [0.004434375, 0.004756007, 0.003911996, 0.0038005, 0.003409971]
         y = [2.48e-188, 7.41e-181, 4.09e-208, 2.08e-223, 2.66e-245]
-        r, p = stats.pearsonr(x,y)
+        r, p = stats.pearsonr(x, y)
 
         # The expected values were computed using mpmath with 80 digits
         # of precision.
@@ -511,35 +513,33 @@ def test_complex_data(self):
         with pytest.raises(ValueError, match=message):
             stats.pearsonr(x, y)
 
-    @pytest.mark.xslow
     @pytest.mark.parametrize('alternative', ('less', 'greater', 'two-sided'))
     @pytest.mark.parametrize('method', ('permutation', 'monte_carlo'))
     def test_resampling_pvalue(self, method, alternative):
         rng = np.random.default_rng(24623935790378923)
-        size = 100 if method == 'permutation' else 1000
+        size = (2, 100) if method == 'permutation' else (2, 1000)
         x = rng.normal(size=size)
         y = rng.normal(size=size)
         methods = {'permutation': stats.PermutationMethod(random_state=rng),
                    'monte_carlo': stats.MonteCarloMethod(rvs=(rng.normal,)*2)}
         method = methods[method]
-        res = stats.pearsonr(x, y, alternative=alternative, method=method)
-        ref = stats.pearsonr(x, y, alternative=alternative)
+        res = stats.pearsonr(x, y, alternative=alternative, method=method, axis=-1)
+        ref = stats.pearsonr(x, y, alternative=alternative, axis=-1)
         assert_allclose(res.statistic, ref.statistic, rtol=1e-15)
         assert_allclose(res.pvalue, ref.pvalue, rtol=1e-2, atol=1e-3)
 
-    @pytest.mark.xslow
     @pytest.mark.parametrize('alternative', ('less', 'greater', 'two-sided'))
     def test_bootstrap_ci(self, alternative):
-        rng = np.random.default_rng(24623935790378923)
-        x = rng.normal(size=100)
-        y = rng.normal(size=100)
-        res = stats.pearsonr(x, y, alternative=alternative)
+        rng = np.random.default_rng(2462935790378923)
+        x = rng.normal(size=(2, 100))
+        y = rng.normal(size=(2, 100))
+        res = stats.pearsonr(x, y, alternative=alternative, axis=-1)
 
         method = stats.BootstrapMethod(random_state=rng)
         res_ci = res.confidence_interval(method=method)
         ref_ci = res.confidence_interval()
 
-        assert_allclose(res_ci, ref_ci, atol=1e-2)
+        assert_allclose(res_ci, ref_ci, atol=1.5e-2)
 
     def test_invalid_method(self):
         message = "`method` must be an instance of..."
@@ -550,6 +550,105 @@ def test_invalid_method(self):
         with pytest.raises(ValueError, match=message):
             res.confidence_interval(method="exact")
 
+    @pytest.mark.parametrize('axis', [0, 1])
+    def test_axis01(self, axis):
+        rng = np.random.default_rng(38572345825)
+        shape = (9, 10)
+        x, y = rng.normal(size=(2,) + shape)
+        res = stats.pearsonr(x, y, axis=axis)
+        ci = res.confidence_interval()
+        if axis == 0:
+            x, y = x.T, y.T
+        for i in range(x.shape[0]):
+            res_i = stats.pearsonr(x[i], y[i])
+            ci_i = res_i.confidence_interval()
+            assert_allclose(res.statistic[i], res_i.statistic)
+            assert_allclose(res.pvalue[i], res_i.pvalue)
+            assert_allclose(ci.low[i], ci_i.low)
+            assert_allclose(ci.high[i], ci_i.high)
+
+    def test_axis_None(self):
+        rng = np.random.default_rng(38572345825)
+        shape = (9, 10)
+        x, y = rng.normal(size=(2,) + shape)
+        res = stats.pearsonr(x, y, axis=None)
+        ci = res.confidence_interval()
+        ref = stats.pearsonr(x.ravel(), y.ravel())
+        ci_ref = ref.confidence_interval()
+        assert_allclose(res.statistic, ref.statistic)
+        assert_allclose(res.pvalue, ref.pvalue)
+        assert_allclose(ci, ci_ref)
+
+    def test_nd_input_validation(self):
+        x, y = np.ones((2, 2, 5))
+        message = '`axis` must be an integer.'
+        with pytest.raises(ValueError, match=message):
+            stats.pearsonr(x, y, axis=1.5)
+
+        message = '`x` and `y` must have the same length along `axis`'
+        with pytest.raises(ValueError, match=message):
+            stats.pearsonr(x, np.ones((2, 1)), axis=1)
+
+        message = '`x` and `y` must have length at least 2.'
+        with pytest.raises(ValueError, match=message):
+            stats.pearsonr(*np.ones((2, 2, 1)), axis=1)
+
+        message = '`x` and `y` must be broadcastable.'
+        with pytest.raises(ValueError, match=message):
+            stats.pearsonr(x, np.ones((3, 5)), axis=1)
+
+    def test_nd_special_cases(self):
+        rng = np.random.default_rng(34989235492245)
+        x0, y0 = rng.random((2, 3, 5))
+
+        message = 'An input array is constant'
+        with pytest.warns(stats.ConstantInputWarning, match=message):
+            x = x0.copy()
+            x[0] = 1
+            res = stats.pearsonr(x, y0, axis=1)
+            ci = res.confidence_interval()
+            np.testing.assert_equal(res.statistic[0], np.nan)
+            np.testing.assert_equal(res.pvalue[0], np.nan)
+            np.testing.assert_equal(ci.low[0], np.nan)
+            np.testing.assert_equal(ci.high[0], np.nan)
+            assert not np.any(np.isnan(res.statistic[1:]))
+            assert not np.any(np.isnan(res.pvalue[1:]))
+            assert not np.any(np.isnan(ci.low[1:]))
+            assert not np.any(np.isnan(ci.high[1:]))
+
+        message = 'An input array is nearly constant'
+        with pytest.warns(stats.NearConstantInputWarning, match=message):
+            x[0, 0] = 1 + 1e-15
+            stats.pearsonr(x, y0, axis=1)
+
+        # length 2 along axis
+        x = [[1, 2], [1, 2], [2, 1], [2, 1]]
+        y = [[1, 2], [2, 1], [1, 2], [2, 1]]
+        res = stats.pearsonr(x, y, axis=-1)
+        ci = res.confidence_interval()
+        assert_allclose(res.statistic, [1, -1, -1, 1])
+        assert_allclose(res.pvalue, 1)
+        assert_allclose(ci.low, -1)
+        assert_allclose(ci.high, 1)
+        assert (res.pvalue.shape == res.statistic.shape
+                == ci.low.shape == ci.high.shape)
+
+    @array_api_compatible
+    @pytest.mark.parametrize('axis', [0, 1, None])
+    @pytest.mark.parametrize('alternative', ['less', 'greater', 'two-sided'])
+    def test_array_api(self, xp, axis, alternative):
+        x, y = rng.normal(size=(2, 10, 11))
+        res = stats.pearsonr(xp.asarray(x), xp.asarray(y),
+                             axis=axis, alternative=alternative)
+        ref = stats.pearsonr(x, y, axis=axis, alternative=alternative)
+        xp_assert_close(res.statistic, xp.asarray(ref.statistic))
+        xp_assert_close(res.pvalue, xp.asarray(ref.pvalue))
+
+        res_ci = res.confidence_interval()
+        ref_ci = ref.confidence_interval()
+        xp_assert_close(res_ci.low, xp.asarray(ref_ci.low))
+        xp_assert_close(res_ci.high, xp.asarray(ref_ci.high))
+
 
 class TestFisherExact:
     """Some tests to show that fisher_exact() works correctly.


Code review comments:
diff_hunk: @@ -4347,47 +4350,47 @@ def _pearsonr_fisher_ci(r, n, confidence_level, alternative):
     Fisher's transformation is used to compute the confidence interval
     (https://en.wikipedia.org/wiki/Fisher_transformation).
     """
-    if r == 1:
-        zr = np.inf
-    elif r == -1:
-        zr = -np.inf
-    else:
-        zr = np.arctanh(r)
+    xp = array_namespace(r)
+
+    with np.errstate(divide='ignore'):
+        zr = xp.arctanh(r)
path: scipy/stats/_stats_py.py
text: @user1:
```suggestion
        zr = xp.atanh(r)
```


@author:
Fixed with https://github.com/scipy/scipy/pull/20284#discussion_r1533596888 locally.

diff_hunk: @@ -1250,6 +1250,12 @@ def boxcox_normmax(
     6.000...
     """
     x = np.asarray(x)
+
+    if np.any(~np.isfinite(x) | (x < 0)):
path: scipy/stats/_morestats.py
text: @user1:
is the `~` idiomatic here? Fine if so, but I haven't seen it before

@author:
[Not uncommon](https://stackoverflow.com/questions/13728708/inverting-a-numpy-boolean-array-using), I don't think. What do you prefer? I used to use `np.logical_not`, but this is a lot shorter.

@user1:
True, that is a lot less verbose, fine with me

@author:
I could rephrase as `not np.all(..)` and it would probably be a bit more efficient, avoiding an array operation.

diff_hunk: @@ -4716,106 +4749,144 @@ def pearsonr(x, y, *, alternative='two-sided', method=None):
     than zero which happens in about half of the cases if we sample x and y.
 
     """
-    n = len(x)
-    if n != len(y):
-        raise ValueError('x and y must have the same length.')
+    xp = array_namespace(x, y)
+    x = xp.asarray(x)
+    y = xp.asarray(y)
+
+    if not is_numpy(xp) and method is not None:
+        method = 'invalid'
+
+    if axis is None:
+        x = xp.reshape(x, (-1,))
+        y = xp.reshape(y, (-1,))
+        axis = -1
+
+    axis_int = int(axis)
+    if axis_int != axis:
+        raise ValueError('`axis` must be an integer.')
+    axis = axis_int
+
+    n = x.shape[axis]
+    if n != y.shape[axis]:
+        raise ValueError('`x` and `y` must have the same length along `axis`.')
 
     if n < 2:
-        raise ValueError('x and y must have length at least 2.')
+        raise ValueError('`x` and `y` must have length at least 2.')
 
-    x = np.asarray(x)
-    y = np.asarray(y)
+    try:
+        x, y = xp.broadcast_arrays(x, y)
+    except ValueError as e:
+        message = '`x` and `y` must be broadcastable.'
+        raise ValueError(message) from e
+
+    x = xp.moveaxis(x, axis, -1)
+    y = xp.moveaxis(y, axis, -1)
+    axis = -1
+
+    dtype = xp.result_type(x.dtype, y.dtype)
+    if xp.isdtype(dtype, "integral"):
+        dtype = xp.float64
 
-    if (np.issubdtype(x.dtype, np.complexfloating)
-            or np.issubdtype(y.dtype, np.complexfloating)):
+    if xp.isdtype(dtype, "complex floating"):
         raise ValueError('This function does not support complex data')
 
+    x = xp.astype(x, dtype, copy=False)
+    y = xp.astype(y, dtype, copy=False)
+    threshold = xp.finfo(dtype).eps ** 0.75
+
     # If an input is constant, the correlation coefficient is not defined.
-    if (x == x[0]).all() or (y == y[0]).all():
+    const_x = xp.all(x == x[..., [0]], axis=-1)
+    const_y = xp.all(y == y[..., [0]], axis=-1)
+    const_xy = const_x | const_y
+    if xp.any(const_xy):
         msg = ("An input array is constant; the correlation coefficient "
                "is not defined.")
         warnings.warn(stats.ConstantInputWarning(msg), stacklevel=2)
-        result = PearsonRResult(statistic=np.nan, pvalue=np.nan, n=n,
-                                alternative=alternative, x=x, y=y)
-        return result
 
     if isinstance(method, PermutationMethod):
-        def statistic(y):
-            statistic, _ = pearsonr(x, y, alternative=alternative)
+        def statistic(y, axis):
+            statistic, _ = pearsonr(x, y, axis=axis, alternative=alternative)
             return statistic
 
         res = permutation_test((y,), statistic, permutation_type='pairings',
-                               alternative=alternative, **method._asdict())
+                               axis=axis, alternative=alternative, **method._asdict())
 
         return PearsonRResult(statistic=res.statistic, pvalue=res.pvalue, n=n,
-                              alternative=alternative, x=x, y=y)
+                              alternative=alternative, x=x, y=y, axis=axis)
     elif isinstance(method, MonteCarloMethod):
-        def statistic(x, y):
-            statistic, _ = pearsonr(x, y, alternative=alternative)
+        def statistic(x, y, axis):
+            statistic, _ = pearsonr(x, y, axis=axis, alternative=alternative)
             return statistic
 
         if method.rvs is None:
             rng = np.random.default_rng()
             method.rvs = rng.normal, rng.normal
 
-        res = monte_carlo_test((x, y,), statistic=statistic,
+        res = monte_carlo_test((x, y,), statistic=statistic, axis=axis,
                                alternative=alternative, **method._asdict())
 
         return PearsonRResult(statistic=res.statistic, pvalue=res.pvalue, n=n,
-                              alternative=alternative, x=x, y=y)
+                              alternative=alternative, x=x, y=y, axis=axis)
+    elif method == 'invalid':
+        message = ('`method` must be None if arguments are not NumPy arrays.')
+        raise ValueError(message)
     elif method is not None:
         message = ('`method` must be an instance of `PermutationMethod`,'
                    '`MonteCarloMethod`, or None.')
         raise ValueError(message)
 
-    # dtype is the data type for the calculations.  This expression ensures
-    # that the data type is at least 64 bit floating point.  It might have
-    # more precision if the input is, for example, np.longdouble.
-    dtype = type(1.0 + x[0] + y[0])
-
     if n == 2:
-        r = dtype(np.sign(x[1] - x[0])*np.sign(y[1] - y[0]))
-        result = PearsonRResult(statistic=r, pvalue=1.0, n=n,
-                                alternative=alternative, x=x, y=y)
+        r = (xp.sign(x[..., 1] - x[..., 0])*xp.sign(y[..., 1] - y[..., 0]))[()]
+        result = PearsonRResult(statistic=r, pvalue=xp.ones_like(r)[()], n=n,
+                                alternative=alternative, x=x, y=y, axis=axis)
         return result
 
-    xmean = x.mean(dtype=dtype)
-    ymean = y.mean(dtype=dtype)
-
-    # By using `astype(dtype)`, we ensure that the intermediate calculations
-    # use at least 64 bit floating point.
-    xm = x.astype(dtype) - xmean
-    ym = y.astype(dtype) - ymean
+    xmean = x.mean(axis=axis, keepdims=True)
+    ymean = y.mean(axis=axis, keepdims=True)
path: scipy/stats/_stats_py.py
text: @user1:
```suggestion
    xmean = xp.mean(x, axis=axis, keepdims=True)
    ymean = xp.mean(y, axis=axis, keepdims=True)
```


@author:
Oops I couldn't just apply this suggestion after the merge for some reason, so I did it locally. Sorry about the force pushes : / I screwed up the merge the first time.