language: python
python:
  - "3.6"
  - "3.7"
  - "3.8"
install:
  - pip install --upgrade pip
  - pip install numpy jaxlib jax tensorflow tensorflow-datasets --upgrade
  - pip install -e .
script:
  - set -e
  - for f in tests/*.py; do python $f; done

language: python
python:
  - "3.6"
  - "3.7"
  - "3.8"
install:
  - pip install --upgrade pip
  - pip install numpy jaxlib tensorflow tensorflow-datasets --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


Use JAX at github head instead of pypi, to fix failing tests
Use JAX at github head instead of pypi, to fix failing tests

PiperOrigin-RevId: 327281929
