jax

22d14fd7 - Remove workaround for Mac linear algebra bug that is fixed in the minimum jaxlib version. (#3080)