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

Code before:
        message = ("The power mean is only defined if all elements are "
                   "non-negative; otherwise, the result is NaN.")
        warnings.warn(message, RuntimeWarning, stacklevel=2)

    with np.errstate(divide='ignore', invalid='ignore'):
        return _xp_mean(a**float(p), axis=axis, weights=weights)**(1/p)


ModeResult = namedtuple('ModeResult', ('mode', 'count'))


def _mode_result(mode, count):
    # When a slice is empty, `_axis_nan_policy` automatically produces
    # NaN for `mode` and `count`. This is a reasonable convention for `mode`,
    # but `count` should not be NaN; it should be zero.
    i = np.isnan(count)
    if i.shape == ():
        count = np.asarray(0, dtype=count.dtype)[()] if i else count
    else:
        count[i] = 0
    return ModeResult(mode, count)


@xp_capabilities(np_only=True)
@_axis_nan_policy_factory(_mode_result, override={'nan_propagation': False})
def mode(a, axis=0, nan_policy='propagate', keepdims=False):
    r"""Return an array of the modal (most common) value in the passed array.

    If there is more than one such value, only one is returned.
    The bin-count for the modal bins is also returned.

    Parameters
    ----------
    a : array_like
        Numeric, n-dimensional array of which to find mode(s).
    axis : int or None, optional
        Axis along which to operate. Default is 0. If None, compute over
        the whole array `a`.
    nan_policy : {'propagate', 'raise', 'omit'}, optional
        Defines how to handle when input contains nan.
        The following options are available (default is 'propagate'):

          * 'propagate': treats nan as it would treat any other value
          * 'raise': throws an error
          * 'omit': performs the calculations ignoring nan values
    keepdims : bool, optional
        If set to ``False``, the `axis` over which the statistic is taken
        is consumed (eliminated from the output array). If set to ``True``,
        the `axis` is retained with size one, and the result will broadcast
        correctly against the input array.


Review comment (suggestion):
Previously, I had `jax_jit=False`, so I expected the arrays to be eager. However, it seems they were not, because it was triggering this code:

https://github.com/scipy/scipy/blob/434fc553a403eea11357bfeb0febb66b5175271e/scipy/stats/_axis_nan_policy.py#L439-L444

(The way that `temp` is determined in general is not ideal, since it basically selects only the first (zeroth) argument of the function to check for laziness. But that is exactly right for `mode`, so I don't think that's the problem.)

Why is the array lazy? JAX arrays should be fine when they're eager - they were working back in the original incarnation of gh-22857 before we made them lazy by default.

Code after:
        message = ("The power mean is only defined if all elements are "
                   "non-negative; otherwise, the result is NaN.")
        warnings.warn(message, RuntimeWarning, stacklevel=2)

    with np.errstate(divide='ignore', invalid='ignore'):
        return _xp_mean(a**float(p), axis=axis, weights=weights)**(1/p)


ModeResult = namedtuple('ModeResult', ('mode', 'count'))


def _mode_result(mode, count):
    # When a slice is empty, `_axis_nan_policy` automatically produces
    # NaN for `mode` and `count`. This is a reasonable convention for `mode`,
    # but `count` should not be NaN; it should be zero.
    xp = array_namespace(mode, count)
    i = xp.isnan(count)
    if i.shape == ():
        count = xp.asarray(0, dtype=count.dtype)[()] if i else count
    else:
        count = xpx.at(count)[i].set(0)
    return ModeResult(mode, count)


@xp_capabilities(skip_backends=[('dask.array', "can't compute chunk size"),
                                ('jax.numpy', "relies on _axis_nan_policy"),
                                ('cupy', "data-apis/array-api-compat#312")])
@_axis_nan_policy_factory(_mode_result, override={'nan_propagation': False})
def mode(a, axis=0, nan_policy='propagate', keepdims=False):
    r"""Return an array of the modal (most common) value in the passed array.

    If there is more than one such value, only one is returned.
    The bin-count for the modal bins is also returned.

    Parameters
    ----------
    a : array_like
        Numeric, n-dimensional array of which to find mode(s).
    axis : int or None, optional
        Axis along which to operate. Default is 0. If None, compute over
        the whole array `a`.
    nan_policy : {'propagate', 'raise', 'omit'}, optional
        Defines how to handle when input contains nan.
        The following options are available (default is 'propagate'):

          * 'propagate': treats nan as it would treat any other value
          * 'raise': throws an error
          * 'omit': performs the calculations ignoring nan values
    keepdims : bool, optional
        If set to ``False``, the `axis` over which the statistic is taken
        is consumed (eliminated from the output array). If set to ``True``,