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

Code before:
    """
    # _axis_nan_policy decorator does not currently support these for lazy arrays.
    # We want to run tests with lazy backends, so don't pass the arguments explicitly
    # unless necessary.
    kwargs = {}
    if keepdims is not False:
        kwargs['keepdims'] = keepdims
    if nan_policy != 'propagate':
        kwargs['nan_policy'] = nan_policy
    res = _yeojohnson_llf(data, lmb=lmb, axis=axis, **kwargs)
    return res[()] if res.ndim == 0 else res


@_axis_nan_policy_factory(lambda x: x, n_outputs=1, default_axis=0,
                          result_to_tuple=lambda x, _: (x,))
def _yeojohnson_llf(data, *, lmb, axis=0):
    xp = array_namespace(data)
    dtype = xp_result_type(lmb, data, force_floating=True, xp=xp)
    data = xp.asarray(data, dtype=dtype)

    n = data.shape[axis]
    if n == 0:
        return _get_nan(data, xp=xp)
    eps = xp.finfo(dtype).eps
    pos = data >= 0  # binary mask

    # There exists numerical instability when abs(lmb) or abs(lmb - 2) is very small
    if xp.all(pos):
        if abs(lmb) < eps:
            logvar = xp.log(xp.var(xp.log1p(data), axis=axis))
        else:
            logvar = _log_var(lmb * xp.log1p(data), xp, axis) - 2 * math.log(abs(lmb))

    elif xp.all(~pos):
        if abs(lmb - 2) < eps:
            logvar = xp.log(xp.var(xp.log1p(-data), axis=axis))
        else:
            logvar = _log_var((2 - lmb) * xp.log1p(-data), xp, axis) - 2 * math.log(
                abs(2 - lmb)
            )

    else:  # mixed positive and negative data
        y = _yeojohnson_transform(data, lmb, xp=xp)
        sigma = xp.var(y, axis=axis)

        # Suppress RuntimeWarning raised by np.log when the variance is too low
        finite_variance = sigma >= xp.finfo(sigma.dtype).smallest_normal
        logvar = xpx.apply_where(finite_variance, (sigma,), xp.log, fill_value=-xp.inf)

    loglike = (-n / 2 * logvar
               + (lmb - 1) * xp.sum(xp.sign(data) * xp.log1p(xp.abs(data)), axis=axis))

Review comment (bug):
```suggestion
    # Special case all-positive/negative data to avoid overflow and precision loss
    pos = data >= 0  # binary mask
```

Code after:
    """
    # _axis_nan_policy decorator does not currently support these for lazy arrays.
    # We want to run tests with lazy backends, so don't pass the arguments explicitly
    # unless necessary.
    kwargs = {}
    if keepdims is not False:
        kwargs['keepdims'] = keepdims
    if nan_policy != 'propagate':
        kwargs['nan_policy'] = nan_policy
    res = _yeojohnson_llf(data, lmb=lmb, axis=axis, **kwargs)
    return res[()] if res.ndim == 0 else res


@_axis_nan_policy_factory(lambda x: x, n_outputs=1, default_axis=0,
                          result_to_tuple=lambda x, _: (x,))
def _yeojohnson_llf(data, *, lmb, axis=0):
    xp = array_namespace(data)
    dtype = xp_result_type(lmb, data, force_floating=True, xp=xp)
    data = xp.asarray(data, dtype=dtype)

    n = data.shape[axis]
    if n == 0:
        return _get_nan(data, xp=xp)
    eps = xp.finfo(dtype).eps
    # Special case all-positive/negative data to avoid overflow and precision loss
    pos = data >= 0  # binary mask

    # There exists numerical instability when abs(lmb) or abs(lmb - 2) is very small
    if xp.all(pos):
        if abs(lmb) < eps:
            logvar = xp.log(xp.var(xp.log1p(data), axis=axis))
        else:
            logvar = _log_var(lmb * xp.log1p(data), xp, axis) - 2 * math.log(abs(lmb))

    elif xp.all(~pos):
        if abs(lmb - 2) < eps:
            logvar = xp.log(xp.var(xp.log1p(-data), axis=axis))
        else:
            logvar = _log_var((2 - lmb) * xp.log1p(-data), xp, axis) - 2 * math.log(
                abs(2 - lmb)
            )

    # overflow/precision loss not reported for mixed data; calculate `logvar` directly
    else:  # mixed positive and negative data
        y = _yeojohnson_transform(data, lmb, xp=xp)
        sigma = xp.var(y, axis=axis)

        # Suppress RuntimeWarning raised by np.log when the variance is too low
        finite_variance = sigma >= xp.finfo(sigma.dtype).smallest_normal
        logvar = xpx.apply_where(finite_variance, (sigma,), xp.log, fill_value=-xp.inf)