[functorch] GHA docs build; nbsphinx -> myst-nb (pytorch/functorch#517)

functorch/.github/workflows/docs.yml

@@ -0,0 +1,39 @@
+name: Build Docs
+on:
+  pull_request:
+    types: [opened, synchronize, reopened]
+  push:
+    branches:
+      - main
+
+jobs:
+
+  build-docs:
+    runs-on: ubuntu-18.04
+    steps:
+      - name: Setup Python
+        uses: actions/setup-python@v2
+        with:
+          python-version: "3.9"
+          architecture: x64
+      - name: Checkout functorch
+        uses: actions/checkout@v2
+      - name: Install PyTorch Nightly
+        run: |
+          python3 -mpip install --pre torch>=1.12.0.dev -f https://download.pytorch.org/whl/nightly/cpu/torch_nightly.html
+      - name: Install functorch
+        run: |
+          python3 setup.py install
+      - name: Install docs requirements
+        run: |
+          cd docs
+          python3 -mpip install -r requirements.txt
+      - name: Build docs
+        run: |
+          cd docs
+          make html
+      - name: Upload docs as GHA artifact
+        uses: actions/upload-artifact@v2
+        with:
+          name: docs.whl
+          path: docs/build/html

functorch/docs/Makefile

@@ -2,7 +2,7 @@
 #
 
 # You can set these variables from the command line.
-SPHINXOPTS    ?= -j auto -WT --keep-going
+SPHINXOPTS    ?= -WT --keep-going
 SPHINXBUILD   ?= sphinx-build
 SPHINXPROJ    ?= functorch
 SOURCEDIR     ?= source

functorch/docs/requirements.txt

@@ -3,7 +3,5 @@ docutils==0.16
 -e git+https://github.com/pytorch/pytorch_sphinx_theme.git#egg=pytorch_sphinx_theme
 sphinxcontrib.katex
 sphinx_copybutton>=0.3.1
-nbsphinx
 IPython
-# Required for nbsphinx: I don't think these can be installed via pip
+myst-nb==0.13.2
-# conda install -c conda-forge pandoc

functorch/docs/source/conf.py

@@ -46,7 +46,7 @@ extensions = [
     # 'sphinxcontrib.katex',
     'sphinx.ext.autosectionlabel',
     'sphinx_copybutton',
-    'nbsphinx',
+    'myst_nb',
 ]
 
 # sys.path.insert(0, os.path.abspath('./notebooks'))
@@ -60,6 +60,9 @@ numpydoc_show_class_members = False
 # duplicated section names that are in different documents.
 autosectionlabel_prefix_document = True
 
+# tell myst to not execute ipynb tutorials.
+jupyter_execute_notebooks = "off"
+
 # katex options
 #
 #

functorch/docs/source/ensembling.ipynb

@@ -1 +0,0 @@
-../../notebooks/ensembling.ipynb

functorch/docs/source/functorch.rst

@@ -9,14 +9,14 @@ Function Transforms
     :toctree: generated
     :nosignatures:
 
+    vmap
     grad
     grad_and_value
-    hessian
+    vjp
+    jvp
     jacrev
     jacfwd
-    jvp
+    hessian
-    vmap
-    vjp
 
 Utilities for working with torch.nn.Modules
 -------------------------------------------

functorch/docs/source/index.rst

@@ -10,7 +10,10 @@ functorch is `JAX-like <https://github.com/google/jax>`_ composable function tra
 It aims to provide composable vmap and grad transforms that work with PyTorch modules
 and PyTorch autograd with good eager-mode performance.
 
-**This library is currently under heavy development - if you have suggestions on the API or use-cases you'd like to be covered, please open an github issue or reach out. We'd love to hear about how you're using the library.**
+.. note::
+   This library is currently under heavy development - if you have suggestions on the API or
+   use-cases you'd like to be covered, please open an github issue or reach out.
+   We'd love to hear about how you're using the library.
 
 Why composable function transforms?
 -----------------------------------
@@ -37,6 +40,12 @@ For a whirlwind tour of how to use the transforms, please check out `this sectio
    :caption: Content
 
    install
-   tutorials
    functorch
 
+.. toctree::
+   :maxdepth: 1
+   :caption: Tutorials
+
+   notebooks/jacobians_hessians.ipynb
+   notebooks/ensembling.ipynb
+   notebooks/per_sample_grads.ipynb

functorch/docs/source/install.rst

@@ -1,5 +1,5 @@
-Install
+Install functorch
-=========
+=================
 
 There are two ways to install functorch:
 

functorch/docs/source/jacobians_hessians.ipynb

@@ -1 +0,0 @@
-../../notebooks/jacobians_hessians.ipynb

functorch/docs/source/notebooks

@@ -0,0 +1 @@
+../../notebooks/

functorch/docs/source/per_sample_grads.ipynb

@@ -1 +0,0 @@
-../../notebooks/per_sample_grads.ipynb

functorch/docs/source/tutorials.rst

@@ -1,10 +0,0 @@
-####################
-  Tutorials
-####################
-
-.. toctree::
-   :maxdepth: 2
-
-   ensembling
-   per_sample_grads
-   jacobians_hessians