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

Code before:
        )

    if a.shape[-1] == 0:  #  Only for *testing* _axis_nan_policy decorator
        return _get_nan(a)

    # for simplicity, ensure that shape is (batch size, sample size) and at least 2d
    shape = a.shape  # remember original shape, though, to ensure correct output shape
    # `reshape` can't infer shape element -1 when array is size 0, so take `prod`
    a = np.reshape(a, (math.prod(shape[:-1]), shape[-1]))
    weights = weights if weights is None else np.reshape(weights, (-1, shape[-1]))

    # This is the empirical equivalent of Eq. (13) with identification
    # function from Table 9 (omitting a factor of 2) in [2] (their y is our
    # data a, their x is our t)
    def first_order(t, i):
        ai = a[i].T
        wi = weights if weights is None else weights[i].T
        return _xp_mean(np.abs((ai <= t) - alpha) * (t - ai), weights=wi, axis=0)

    # axis is -1 per the decorator. Use mean to produce NaNs of the correct shape/type.
    x0 = np.min(a, axis=-1)
    x1 = np.max(a, axis=-1)

    # Note that the expectile is the unique solution, so no worries about
    # finding a wrong root.
    i = np.arange(a.shape[0])
    res = find_root(first_order, (x0, x1), args=(i,))

    return np.reshape(res.x, shape[:-1])[()]


def _lmoment_iv(sample, order, axis, sorted, standardize, xp):
    # input validation/standardization for `lmoment`
    sample = xp_promote(sample, force_floating=True, xp=xp)

    message = "`sample` must be an array of real numbers."
    if not xp.isdtype(sample.dtype, "real floating"):
        raise ValueError(message)

    message = "`order` must be a scalar or a non-empty array of positive integers."
    order = xp.arange(1, 5) if order is None else xp.asarray(order)
    if (not xp.isdtype(order.dtype, "integral") or xp.any(order <= 0)
            or order.size == 0 or order.ndim > 1):
        raise ValueError(message)

    # input validation of non-array types can still be performed with NumPy
    axis = np.asarray(axis)[()]
    message = "`axis` must be an integer."
    if not np.issubdtype(axis.dtype, np.integer) or axis.ndim != 0:
        raise ValueError(message)
    axis = int(axis)

Review comment (suggestion):
`i` is the fancy indices of the slices of `a`(/`weights`) for which root finding has not converged. Initially, it's all of them, but some slices may need fewer iterations than others, and those slices (and the corresponding elements of `i`) disappear to avoid doing unnecessary work.

Code after:
    a, weights = xp_promote(a, weights, force_floating=True, xp=xp)

    if alpha < 0 or alpha > 1:
        raise ValueError(
            "The expectile level alpha must be in the range [0, 1]."
        )

    if a.shape[-1] == 0:  #  Only for *testing* _axis_nan_policy decorator
        return _get_nan(a, xp=xp)

    # for simplicity, ensure that shape is (batch size, sample size) and at least 2d
    shape = a.shape  # remember original shape, though, to ensure correct output shape
    # `reshape` can't infer shape element -1 when array is size 0, so take `prod`
    a = xp.reshape(a, (math.prod(shape[:-1]), shape[-1]))
    weights = weights if weights is None else xp.reshape(weights, (-1, shape[-1]))

    # This is the empirical equivalent of Eq. (13) with identification
    # function from Table 9 (omitting a factor of 2) in [2] (their y is our
    # data a, their x is our t)
    def first_order(t, i):
        ai = a[i, ...].T
        wi = weights if weights is None else weights[i, ...].T
        return _xp_mean(xp.abs(xp.astype(ai <= t, ai.dtype) - alpha) * (t - ai),
                        weights=wi, axis=0)

    # axis is -1 per the decorator. Use mean to produce NaNs of the correct shape/type.
    x0 = xp.min(a, axis=-1)
    x1 = xp.max(a, axis=-1)

    # Note that the expectile is the unique solution, so no worries about
    # finding a wrong root.
    i = xp.arange(a.shape[0])
    res = find_root(first_order, (x0, x1), args=(i,))
    res = xp.reshape(res.x, shape[:-1])
    return res[()] if res.ndim == 0 else res


def _lmoment_iv(sample, order, axis, sorted, standardize, xp):
    # input validation/standardization for `lmoment`
    sample = xp_promote(sample, force_floating=True, xp=xp)

    message = "`sample` must be an array of real numbers."
    if not xp.isdtype(sample.dtype, "real floating"):
        raise ValueError(message)

    message = "`order` must be a scalar or a non-empty array of positive integers."
    order = xp.arange(1, 5) if order is None else xp.asarray(order)
    if (not xp.isdtype(order.dtype, "integral") or xp.any(order <= 0)
            or order.size == 0 or order.ndim > 1):
        raise ValueError(message)