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

Code before:
    n = x.shape[-1]

    # "Rearrange the data as (X(1), Y(1)), . . . ,(X(n), Y(n))
    # such that X(1) ≤ ··· ≤ X(n)"
    j = np.argsort(x, axis=-1)
    j, y = np.broadcast_arrays(j, y)
    y = np.take_along_axis(y, j, axis=-1)

    # "Let ri be the rank of Y(i), that is, the number of j such that Y(j) ≤ Y(i)"
    r = stats.rankdata(y, method='max', axis=-1)
    # " additionally define li to be the number of j such that Y(j) ≥ Y(i)"
    # Could probably compute this from r, but that can be an enhancement
    l = stats.rankdata(-y, method='max', axis=-1)

    num = np.sum(np.abs(np.diff(r, axis=-1)), axis=-1)
    if y_continuous:  # [1] Eq. 1.1
        statistic = 1 - 3 * num / (n ** 2 - 1)
    else:  # [1] Eq. 1.2
        den = 2 * np.sum((n - l) * l, axis=-1)
        statistic = 1 - n * num / den

    return statistic, r, l


def _xi_std(r, l, y_continuous):
    # Compute asymptotic standard deviation of xi under null hypothesis of independence

    # `axis=-1` is guaranteed by _axis_nan_policy decorator
    n = np.float64(r.shape[-1])

    # "Suppose that X and Y are independent and Y is continuous. Then
    # √n·ξn(X, Y) → N(0, 2/5) in distribution as n → ∞"
    if y_continuous:  # [1] Theorem 2.1
        return np.sqrt(2 / 5) / np.sqrt(n)

    # "Suppose that X and Y are independent. Then √n·ξn(X, Y)
    # converges to N(0, τ²) in distribution as n → ∞
    # [1] Eq. 2.2 and surrounding math
    i = np.arange(1, n + 1)
    u = np.sort(r, axis=-1)
    v = np.cumsum(u, axis=-1)
    an = 1 / n**4 * np.sum((2*n - 2*i + 1) * u**2, axis=-1)
    bn = 1 / n**5 * np.sum((v + (n - i)*u)**2, axis=-1)
    cn = 1 / n**3 * np.sum((2*n - 2*i + 1) * u, axis=-1)
    dn = 1 / n**3 * np.sum((l * (n - l)), axis=-1)
    tau2 = (an - 2*bn + cn**2) / dn**2

    return np.sqrt(tau2) / np.sqrt(n)


def _chatterjeexi_iv(y_continuous, method):

Review comment (suggestion):
I thought we used to have to shim around `shape` calls, I must be falling behind.

Code after:
    j = xp.argsort(x, axis=-1)
    j, y = xp.broadcast_arrays(j, y)
    y = xp.take_along_axis(y, j, axis=-1)

    # "Let ri be the rank of Y(i), that is, the number of j such that Y(j) ≤ Y(i)"
    r = stats.rankdata(y, method='max', axis=-1)
    # " additionally define li to be the number of j such that Y(j) ≥ Y(i)"
    # Could probably compute this from r, but that can be an enhancement
    l = stats.rankdata(-y, method='max', axis=-1)
    r, l = xp.astype(r, x.dtype), xp.astype(l, x.dtype)

    num = xp.sum(xp.abs(xp.diff(r, axis=-1)), axis=-1)
    if y_continuous:  # [1] Eq. 1.1
        statistic = 1 - 3 * num / (n ** 2 - 1)
    else:  # [1] Eq. 1.2
        den = 2 * xp.sum((n - l) * l, axis=-1)
        statistic = 1 - n * num / den

    return statistic, r, l


def _xi_std(r, l, y_continuous, xp):
    # Compute asymptotic standard deviation of xi under null hypothesis of independence

    # `axis=-1` is guaranteed by _axis_nan_policy decorator
    n = r.shape[-1]

    # "Suppose that X and Y are independent and Y is continuous. Then
    # √n·ξn(X, Y) → N(0, 2/5) in distribution as n → ∞"
    if y_continuous:  # [1] Theorem 2.1
        return xp.asarray(math.sqrt(2 / 5) / math.sqrt(n), dtype=r.dtype)

    # "Suppose that X and Y are independent. Then √n·ξn(X, Y)
    # converges to N(0, τ²) in distribution as n → ∞
    # [1] Eq. 2.2 and surrounding math
    i = xp.arange(1, n + 1, dtype=r.dtype)
    u = xp.sort(r, axis=-1)
    v = xp.cumulative_sum(u, axis=-1)
    an = 1 / n**4 * xp.sum((2*n - 2*i + 1) * u**2, axis=-1)
    bn = 1 / n**5 * xp.sum((v + (n - i)*u)**2, axis=-1)
    cn = 1 / n**3 * xp.sum((2*n - 2*i + 1) * u, axis=-1)
    dn = 1 / n**3 * xp.sum((l * (n - l)), axis=-1)
    tau2 = (an - 2*bn + cn**2) / dn**2

    return xp.sqrt(tau2) / math.sqrt(n)


def _chatterjeexi_iv(y_continuous, method):
    # Input validation for `chatterjeexi`
    # x, y, `axis` input validation taken care of by decorator