language: python
python:
  - "3.6"
  - "3.7"
  - "3.8"
install:
  - pip install --upgrade pip
  - pip install more-itertools numpy jaxlib optax --upgrade
  - pip install git+https://github.com/google/jax.git
  - pip install -e .
script:
  - set -e
  - for f in tests/*.py; do python $f; done

language: python
python:
  - "3.7"
  - "3.8"
  - "3.9"
install:
  - pip install --upgrade pip
  - pip install more-itertools numpy jaxlib optax --upgrade
  - pip install git+https://github.com/google/jax.git
  - pip install -e .
script:
  - set -e
  - for f in tests/*.py; do python $f; done

Update tests to use python 3.9 instead of python 3.6.
Update tests to use python 3.9 instead of python 3.6.