jax

0080c895 - Fix a few type annotations in `api.py`. (#2387)