jax

f49ab50e - expose lax._safe_mul again (c.f. #608)