jax

1bb0f154 - Fixed jupytext failures at head