jax

126909b6 - Update lax.linalg.lu primitive to use registration helper functions.