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

Code before:
        return gmean(a, axis=axis, dtype=dtype, weights=weights)
    elif math.isinf(p):
        message = "Power mean only implemented for finite `p`"
        raise NotImplementedError(message)

    xp = array_namespace(a, weights)
    dtype = (xp_result_type(a, p, weights, force_floating=True, xp=xp)
             if dtype is None else dtype)
    a = xp.asarray(a, dtype=dtype)

    if weights is not None:
        weights = xp.asarray(weights, dtype=dtype)

    negative_mask = a < 0
    if xp.any(negative_mask):
        # `where` avoids having to be careful about dtypes and will work with
        # JAX. This is the exceptional case, so it's OK to be a little slower.
        # Won't work for array_api_strict for now, but see data-apis/array-api#807
        a = xp.where(negative_mask, np.nan, a)
        message = ("The power mean is only defined if all elements are "
                   "non-negative; otherwise, the result is NaN.")
        warnings.warn(message, RuntimeWarning, stacklevel=2)

    # Linearized approximation for small p to avoid numerical issues; see gh-23407
    p_threshold = 2e-6
    if p <= p_threshold:
        return _linearized_pmean(a, p, axis=axis, weights=weights, xp=xp)

    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")])

Review comment (suggestion):
```suggestion
    if abs(p) <= p_threshold:
```
My mistake!

Code after:
        return gmean(a, axis=axis, dtype=dtype, weights=weights)
    elif math.isinf(p):
        message = "Power mean only implemented for finite `p`"
        raise NotImplementedError(message)

    xp = array_namespace(a, weights)
    dtype = (xp_result_type(a, p, weights, force_floating=True, xp=xp)
             if dtype is None else dtype)
    a = xp.asarray(a, dtype=dtype)

    if weights is not None:
        weights = xp.asarray(weights, dtype=dtype)

    negative_mask = a < 0
    if xp.any(negative_mask):
        # `where` avoids having to be careful about dtypes and will work with
        # JAX. This is the exceptional case, so it's OK to be a little slower.
        # Won't work for array_api_strict for now, but see data-apis/array-api#807
        a = xp.where(negative_mask, np.nan, a)
        message = ("The power mean is only defined if all elements are "
                   "non-negative; otherwise, the result is NaN.")
        warnings.warn(message, RuntimeWarning, stacklevel=2)

    # Linearized approximation for small p to avoid numerical issues; see gh-23407
    p_threshold = 2e-6
    if abs(p) <= p_threshold:
        return _linearized_pmean(a, p, axis=axis, weights=weights, xp=xp)

    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")])