jax: Gradient of QR fails

There seems to be a major bug in the gradient of the QR decomposition. The forward-mode derivative is accurate for Q but very far off for R. There also seems to be no check in the tests making sure the gradient is accurate.

This reproduces the error:

import jax
import jax.numpy as jnp
import numpy as np

x = np.random.randn(3, 3)
dx = np.random.randn(3, 3)

primals, tangents = jax.jvp(jnp.linalg.qr, (x,), (dx,))
q, r = primals
dq, dr = tangents

dt = 1e-6
dq_ = (np.linalg.qr(x + dt * dx)[0] - np.linalg.qr(x)[0]) / dt
dr_ = (np.linalg.qr(x + dt * dx)[1] - np.linalg.qr(x)[1]) / dt

assert jnp.allclose(x, q @ r, atol=1e-5, rtol=1e-5)  # passes
assert jnp.allclose(dq, dq_, atol=1e-5, rtol=1e-5)  # passes
assert jnp.allclose(dx, q @ dr_ + dq_ @ r, atol=1e-5, rtol=1e-5)  # passes
assert jnp.allclose(dr, dr_, atol=1e-5, rtol=1e-5)  # fails