Moved all notebooks to docs/notebooks. (jax-ml#1493)

* Moved all notebooks to docs/notebooks.

Now all notebooks are in the same place, thus all are subject
to auto-doc generation at readthedocs.io and to automated testing
with travis.

Some notebooks are too slow, exclude them at docs/conf.py:exclude_patterns.

Cleanup a bit the section headings in notebooks so that they show
up well in readtehdocs.io.

* Increase the cell timeout for executing notebooks
* Exclude also the neural network notebook from auto-generation (timing out)
* Disable the score_matching notebook from auto-doc (travis does not have sklearn)

Author: gnecula

README.md

@@ -85,16 +85,17 @@ open](https://github.com/google/jax) by a growing number of
 ## Quickstart: Colab in the Cloud
 Jump right in using a notebook in your browser, connected to a Google Cloud GPU. Here are some starter notebooks:
 - [The basics: NumPy on accelerators, `grad` for differentiation, `jit` for compilation, and `vmap` for vectorization](https://colab.research.google.com/github/google/jax/blob/master/docs/notebooks/quickstart.ipynb)
-- [Training a Simple Neural Network, with PyTorch Data Loading](https://colab.research.google.com/github/google/jax/blob/master/notebooks/neural_network_and_data_loading.ipynb)
-- [Training a Simple Neural Network, with TensorFlow Dataset Data Loading](https://colab.research.google.com/github/google/jax/blob/master/notebooks/neural_network_with_tfds_data.ipynb)
+- [Training a Simple Neural Network, with PyTorch Data Loading](https://colab.research.google.com/github/google/jax/blob/master/docs/notebooks/Neural_Network_and_Data_Loading.ipynb)
+- [Training a Simple Neural Network, with TensorFlow Dataset Data Loading](https://colab.research.google.com/github/google/jax/blob/master/docs/notebooks/neural_network_with_tfds_data.ipynb)
 
 And for a deeper dive into JAX:
 - [Common gotchas and sharp edges](https://colab.research.google.com/github/google/jax/blob/master/docs/notebooks/Common_Gotchas_in_JAX.ipynb)
 - [The Autodiff Cookbook, Part 1: easy and powerful automatic differentiation in JAX](https://colab.research.google.com/github/google/jax/blob/master/docs/notebooks/autodiff_cookbook.ipynb)
-- [Directly using XLA in Python](https://colab.research.google.com/github/google/jax/blob/master/notebooks/XLA_in_Python.ipynb)
-- [MAML Tutorial with JAX](https://colab.research.google.com/github/google/jax/blob/master/notebooks/maml.ipynb)
-- [Generative Modeling by Estimating Gradients of Data Distribution in JAX](https://colab.research.google.com/github/google/jax/blob/master/notebooks/score_matching.ipynb).
-[How JAX primitives work](https://colab.research.google.com/github/google/jax/blob/master/docs/notebooks/How_JAX_primitives_work.ipynb)
+- [Directly using XLA in Python](https://colab.research.google.com/github/google/jax/blob/master/docs/notebooks/XLA_in_Python.ipynb)
+- [How JAX primitives work](https://colab.research.google.com/github/google/jax/blob/master/docs/notebooks/How_JAX_primitives_work.ipynb)
+- [MAML Tutorial with JAX](https://colab.research.google.com/github/google/jax/blob/master/docs/notebooks/maml.ipynb)
+- [Generative Modeling by Estimating Gradients of Data Distribution in JAX](https://colab.research.google.com/github/google/jax/blob/master/docs/notebooks/score_matching.ipynb).
+
 
 ## Installation
 JAX is written in pure Python, but it depends on XLA, which needs to be compiled

docs/conf.py

@@ -89,7 +89,13 @@
 # List of patterns, relative to source directory, that match files and
 # directories to ignore when looking for source files.
 # This pattern also affects html_static_path and html_extra_path.
-exclude_patterns = []
+exclude_patterns = [
+    # Slow notebook: long time to load tf.ds
+    'notebooks/neural_network_with_tfds_data.ipynb',
+    'notebooks/Neural_Network_and_Data_Loading.ipynb',
+    # Next one needs sklearn, which is not available at Travis.
+    'notebooks/score_matching.ipynb'
+]
 
 # The name of the Pygments (syntax highlighting) style to use.
 pygments_style = None
@@ -101,6 +107,7 @@
 # -- Options for nbsphinx -----------------------------------------------------
 
 # Execute notebooks before conversion: 'always', 'never', 'auto' (default)
+# We execute all notebooks, exclude the slow ones using 'exclude_patterns'
 nbsphinx_execute = 'always'
 
 # Use this kernel instead of the one stored in the notebook metadata:
@@ -114,7 +121,7 @@
 
 
 # Controls when a cell will time out (defaults to 30; use -1 for no timeout):
-nbsphinx_timeout = 60
+nbsphinx_timeout = 90
 
 # Default Pygments lexer for syntax highlighting in code cells:
 #nbsphinx_codecell_lexer = 'ipython3'

docs/developer.rst

@@ -94,18 +94,35 @@ The Colab notebooks are tested for errors as part of the documentation build.
 Update documentation
 ====================
 
-To rebuild the documentation,
-install the ``sphinx`` and ``sphinx_rtd_theme`` pip
-packages and then run at the top-level:
+To rebuild the documentation, install several packages:
 
 .. code-block:: shell
 
-  sphinx-build -M html docs docs/build
+  pip install -r docs/requirements.txt
 
+You must also install ``pandoc`` in order to regenerate the notebooks.
+See `Install Pandoc <https://pandoc.org/installing.html>`_. On Mac, I had success with
+the miniconda installer, then ``conda install -c conda-forge pandoc``.
+If you do not want to install ``pandoc`` then you should regenerate the documentation
+without the notebooks.
+
+You run at top-level one of the following commands:
+
+.. code-block:: shell
+
+  sphinx-build -b html docs docs/build/html  # with the notebooks
+  sphinx-build -b html -D nbsphinx_execute=never docs docs/build/html  # without the notebooks
 
 You can then see the generated documentation in
 ``docs/build/html/index.html``.
 
+Update notebooks
+----------------
+
+Open the notebook with http://colab.research.google.com, update it, ``Run all cells`` then
+``Download ipynb``. You may want to test that it executes properly, using ``sphinx-build`` as
+explained above.
+
 Documentation building on readthedocs.io
 ----------------------------------------
 
@@ -115,7 +132,7 @@ The documentation building is controlled for the entire project by the
 `readthedocs JAX settings <https://readthedocs.org/dashboard/jax>`_. The current settings
 trigger a documentation build as soon as code is pushed to the GitHub ``master`` branch.
 For each code version, the building process is driven by the 
-:file:`.readthedocs.yml` and the ``docs/conf.py`` configuration files.
+``.readthedocs.yml`` and the ``docs/conf.py`` configuration files.
 
 For each automated documentation build you can see the 
 `documentation build logs <https://readthedocs.org/projects/jax/builds/>`_.

docs/index.rst

@@ -13,13 +13,20 @@ For an introduction to JAX, start at the
 
    notebooks/quickstart
    notebooks/autodiff_cookbook
+   Training a Simple Neural Network, with PyTorch Data Loading <https://github.com/google/jax/blob/master/docs/notebooks/Neural_Network_and_Data_Loading.ipynb>
+
 
 .. toctree::
    :maxdepth: 1
-   :caption: Advanced Tutorials
+   :caption: Advanced JAX Tutorials
 
    notebooks/Common_Gotchas_in_JAX
+   notebooks/XLA_in_Python
    notebooks/How_JAX_primitives_work
+   Training a Simple Neural Network, with Tensorflow Datasets Data Loading <https://github.com/google/jax/blob/master/docs/notebooks/neural_network_with_tfds_data.ipynb>
+   notebooks/maml
+   Generative Modeling by Estimating Gradients of Data Distribution <https://github.com/google/jax/blob/master/docs/notebooks/score_matching.ipynb>
+   notebooks/vmapped_log_probs
 
 .. toctree::
    :maxdepth: 1