Disable omnistaging by default

.travis.yml

@@ -60,7 +60,7 @@ install:
   - conda create --yes -q -n pyenv python=$TRAVIS_PYTHON_VERSION
   - conda activate pyenv
   - conda install --yes -q mkl numpy scipy pip mkl-service graphviz cython libgpuarray pygpu
-  - if [[ "$TRAVIS_PYTHON_VERSION" != "3.6" ]]; then conda install --yes -q -c conda-forge 'jax<0.2.0' 'jaxlib'; fi
+  - if [[ "$TRAVIS_PYTHON_VERSION" != "3.6" ]]; then conda install --yes -q -c conda-forge jax jaxlib; fi
   - pip install -q -r requirements.txt
   - conda list && pip freeze
   - python -c 'import theano; print(theano.config.__str__(print_doc=False))'

requirements.txt

@@ -10,5 +10,5 @@ coveralls
 cython
 sympy
 versioneer
-jax<0.2.0; python_version > '3.6'
+jax; python_version > '3.6'
 jaxlib; python_version > '3.6'

theano/sandbox/jaxify.py

@@ -51,6 +51,10 @@ from theano.compile.ops import (
 from theano.tensor.opt import MakeVector
 
 
+# XXX: Enabling this will break some shape-based functionality, and severely
+# limit the types of graphs that can be converted.
+# See https://github.com/google/jax/blob/4d556837cc9003492f674c012689efc3d68fdf5f/design_notes/omnistaging.md
+jax.config.disable_omnistaging()
 jax.config.update("jax_enable_x64", True)
 
 subtensor_ops = (Subtensor, AdvancedSubtensor1, BaseAdvancedSubtensor)