jax

8358c016 - Merge pull request #4890 from google:autodiff-cookbook-fix