jax

a57a1a39 - increase minimum jaxlib version and remove some janky feature detection