Add utility for scipy.linalg.qr_update

[YigitElma]

For my application, I need to take QR of

$$ \tilde{Q} \tilde{R} = \begin{pmatrix} J \\ \alpha I \end{pmatrix} $$

where I already have

$$ QR = J $$

So, I have tried to implement this myself using Givens rotations. But due to N*M loops it is very slow, which makes it useless since the main goal of this function is to calculate QR faster than otherwise calling jax.scipy.linalg.qr. I guess this would be useful for people using jax for trust region Newton method optimization. I thought of using Householder reflections but I couldn't implement it under jax.jit. For reference, this is implemented in scipy here and calls some LAPACK routines in a for loop. I hope this can be implemented in the same way some other matrix decomposition routines are implemented in JAX.

I naively tried to implement this myself purely in Python like this,

def _givens_jax(a, b):
    b_zero = abs(b) == 0
    a_lt_b = abs(a) < abs(b)
    t = -jnp.where(a_lt_b, a, b) / jnp.where(a_lt_b, b, a)
    r = rsqrt(1 + abs(t) ** 2).astype(t.dtype)
    cs = jnp.where(b_zero, 1, jnp.where(a_lt_b, r * t, r))
    sn = jnp.where(b_zero, 0, jnp.where(a_lt_b, r, r * t))
    G2 = jnp.array([[cs, -sn], [sn, cs]])
    return G2.astype(float)

@jax.jit
def update_qr_jax(A, w, q, r):
    """Update QR factorization with a diagonal matrix w at the bottom."""
    m, n = A.shape
    Q = jnp.eye(m + n)
    Q = Q.at[:m, :m].set(q)

    R = jnp.vstack([r, w])

    def body_inner(i, jQR):
        j, Q, R = jQR
        i = m + j - i
        a, b = R[i - 1, j], R[i, j]
        G2 = _givens_jax(a, b)
        R = R.at[jnp.array([i - 1, i])].set(G2 @ R[jnp.array([i - 1, i])])
        Q = Q.at[:, jnp.array([i - 1, i])].set(Q[:, jnp.array([i - 1, i])] @ G2.T)
        return j, Q, R

    def body(j, QR):
        Q, R = QR
        j, Q, R = fori_loop(0, m, body_inner, (j, Q, R))
        return Q, R

    Q, R = fori_loop(0, n, body, (Q, R))
    R = jnp.where(jnp.abs(R) < 1e-10, 0, R)

    return Q, R

Note: I also tried economic mode QR to reduce matrix size, but this is still slow.