Update JAX setup for READTHEDOCS and travis

Author: lululxvi

.travis.yml

@@ -28,10 +28,12 @@ install:
 
   - conda install -c conda-forge scikit-optimize
   - pip install tensorflow tensorflow-probability
-  - conda install pytorch torchvision torchaudio cpuonly -c pytorch
+  - conda install pytorch cpuonly -c pytorch
+  - pip install jax flax optax
 
 script:
   # Your test script goes here
   - DDEBACKEND=tensorflow.compat.v1 python -c "import deepxde"
   - DDEBACKEND=tensorflow python -c "import deepxde"
   - DDEBACKEND=pytorch python -c "import deepxde"
+  - DDEBACKEND=jax python -c "import deepxde"

deepxde/nn/__init__.py

@@ -23,6 +23,7 @@
 if os.environ.get("READTHEDOCS") == "True":
     # The backend should be tensorflow/tensorflow.compat.v1 to ensure backend.tf is not
     # None.
+    from . import jax
     from . import pytorch
     from . import tensorflow
     from . import tensorflow_compat_v1

docs/index.rst

@@ -82,11 +82,12 @@ If you are looking for information on a specific function, class or method, this
   modules/deepxde
   modules/deepxde.data
   modules/deepxde.geometry
-  modules/deepxde.icbcs
+  modules/deepxde.icbc
   modules/deepxde.nn
-  modules/deepxde.nn.tensorflow_compat_v1
-  modules/deepxde.nn.tensorflow
+  modules/deepxde.nn.jax
   modules/deepxde.nn.pytorch
+  modules/deepxde.nn.tensorflow
+  modules/deepxde.nn.tensorflow_compat_v1
   modules/deepxde.optimizers
   modules/deepxde.utils
 

docs/requirements.txt

@@ -3,7 +3,15 @@ numpy
 scikit-learn
 scikit-optimize
 scipy
+docutils<0.18  # https://github.com/readthedocs/readthedocs.org/issues/8616
+
+# TensorFlow 1.x
 tensorflow>=2.2.0
+# TensorFlow 2.x
 tensorflow-probability>=0.10.0
+# PyTorch
 torch
-docutils<0.18  # https://github.com/readthedocs/readthedocs.org/issues/8616
+# JAX
+jax
+flax
+optax