Loess quantiles

import numba
import numpy as np

def loess_quantiles(x, xp, fp, q, *, bw):
    """Compute quantiles using local kernel approximations
    
    :param q:
        the quantiles to evaluate
    :param bw:
        the scales of the Gaussian kernel
    """
    x, xp, fp, q, bw = (np.asarray(a) for a in (x, xp, fp, q, bw))
    
    assert x.ndim == 2
    assert xp.ndim == 2 and xp.shape[1] == x.shape[1]
    assert fp.ndim == 1 and fp.shape[0] == xp.shape[0]
    assert bw.ndim == 1 and bw.shape[0] == x.shape[1]
    
    res = np.zeros((len(x), len(q)), dtype=fp.dtype)
    sort_idx = np.argsort(fp)
    
    _loess_quantiles(res, x, xp[sort_idx], fp[sort_idx], q, bw)

    return res


@numba.jit(nopython=True, nogil=True)
def _loess_quantiles(res, x, xp, fp, q, bw):
    bw = bw.reshape(1, -1)
    
    for i in range(len(x)):
        xi = x[i, :].reshape(1, -1)
        dx = (xp[:, :] - xi) / bw

        log_w = -0.5 * np.sum(dx ** 2.0, axis=1)
        
        w = np.cumsum(np.exp(log_w - log_w.max()))
        w = w / w[-1]

        res[i, :] = np.interp(q, w, fp)