Merge pull request #1 from google/master

Update branch to head.

Author: sschoenholz

.github/workflows/build_mac_jaxlib.yml

@@ -9,3 +9,8 @@
 build/
 dist/
 .mypy_cache/
+.pytype/
+docs/notebooks/.ipynb_checkpoints/
+docs/_autosummary
+.idea
+.vscode

.gitignore

@@ -0,0 +1,20 @@
+# .readthedocs.yml
+# Read the Docs configuration file
+# See https://docs.readthedocs.io/en/stable/config-file/v2.html for details
+
+# Required
+version: 2
+
+# Build documentation in the docs/ directory with Sphinx
+sphinx:
+  configuration: docs/conf.py
+
+# Optionally build your docs in additional formats such as PDF and ePub
+formats:
+  - htmlzip
+
+# Optionally set the version of Python and requirements required to build your docs
+python:
+  version: 3.7
+  install:
+    - requirements: docs/requirements.txt

.readthedocs.yml

@@ -1,27 +1,51 @@
-sudo: false
 notifications:
   email: false
 dist: xenial
 language: python
-python:
-  - "2.7"
-  - "3.6"
-env:
-  - JAX_ENABLE_X64=0 JAX_NUM_GENERATED_CASES=25
-  - JAX_ENABLE_X64=1 JAX_NUM_GENERATED_CASES=25
+os: linux
+jobs:
+  include:
+    - python: "3.6"
+      env: JAX_ONLY_CHECK_TYPES=true
+    - python: "3.6"
+      env: JAX_ENABLE_X64=0 JAX_NUM_GENERATED_CASES=25
+    - python: "3.6"
+      env: JAX_ENABLE_X64=1 JAX_NUM_GENERATED_CASES=25
+    - python: "3.7"
+      env: JAX_ENABLE_X64=1 JAX_ONLY_DOCUMENTATION=true
+
 before_install:
-  - if [[ "$TRAVIS_PYTHON_VERSION" == "2.7" ]]; then
-      wget https://repo.continuum.io/miniconda/Miniconda2-latest-Linux-x86_64.sh -O miniconda.sh;
-    else
-      wget https://repo.continuum.io/miniconda/Miniconda3-latest-Linux-x86_64.sh -O miniconda.sh;
-    fi
+  - wget https://repo.continuum.io/miniconda/Miniconda3-latest-Linux-x86_64.sh -O miniconda.sh;
   - bash miniconda.sh -b -p $HOME/miniconda
   - export PATH="$HOME/miniconda/bin:$PATH"
-  - conda update --yes conda
+  - conda config --set always_yes yes --set changeps1 no --set show_channel_urls yes
+  - conda config --set channel_priority strict
+  - conda config --set add_pip_as_python_dependency yes
+  - conda config --remove channels defaults
   - conda config --add channels conda-forge
+  - conda update -q conda
 install:
-  - conda install --yes python=$TRAVIS_PYTHON_VERSION pip six protobuf>=3.6.0 absl-py opt_einsum numpy scipy pytest-xdist fastcache
-  - pip install jaxlib
+  - conda install --yes python=$TRAVIS_PYTHON_VERSION pip absl-py opt_einsum numpy scipy pytest-xdist pytest-benchmark mypy
+  - if [ "$JAX_ONLY_CHECK_TYPES" = true ]; then
+      pip install pytype ;
+    fi
+  # The jaxlib version should match the minimum jaxlib version in
+  # jax/lib/__init__.py. This tests JAX PRs against the oldest permitted
+  # jaxlib.
+  - pip install jaxlib==0.1.43
   - pip install -v .
+  # The following are needed to test the Colab notebooks and the documentation building
+  - if [ "$JAX_ONLY_DOCUMENTATION" = true ]; then
+      conda install --yes -c conda-forge pandoc ipykernel;
+      conda install --yes sphinx sphinx_rtd_theme nbsphinx sphinx-autodoc-typehints jupyter_client matplotlib;
+      pip install sklearn;
+    fi
 script:
-  - pytest -n 1 tests examples -W ignore
+  - if [ "$JAX_ONLY_DOCUMENTATION" = true ]; then
+      sphinx-build -b html -D nbsphinx_execute=always docs docs/build/html ;
+    elif [ "$JAX_ONLY_CHECK_TYPES" = true ]; then
+      echo "===== Checking with mypy ====" &&
+      time mypy --config-file=mypy.ini jax ;
+    else
+      pytest tests examples -W ignore ;
+    fi

.travis.yml

@@ -0,0 +1,3 @@
+# Change Log
+
+See [docs/CHANGELOG.rst](https://jax.readthedocs.io/en/latest/CHANGELOG.html).