Merge pull request #526 from google:doc4

PiperOrigin-RevId: 562759212

.github/workflows/tests.yml

@@ -43,3 +43,29 @@ jobs:
           python -c "import jaxlib; print('jaxlib', jaxlib.__version__)"
           pytest tests
         shell: bash
+
+
+  build-and-test-docs:
+    name: "Build documentation"
+    runs-on: ubuntu-latest
+    steps:
+      - name: Cancel previous
+        uses: styfle/cancel-workflow-action@0.11.0
+        with:
+          access_token: ${{ github.token }}
+        if: ${{github.ref != 'refs/heads/main'}}
+      - uses: "actions/checkout@v3"
+      - name: Set up Python 3.11
+        uses: "actions/setup-python@v4"
+        with:
+            python-version: 3.11
+      - name: Install dependencies
+        run: |
+          set -xe
+          pip install --upgrade pip setuptools wheel
+          pip install -r docs/requirements.txt
+      - name: Build documentation
+        run: |
+          set -xe
+          python -VV
+          cd docs && make clean && make html

docs/requirements.txt

@@ -11,6 +11,8 @@ sphinx-remove-toctrees>=0.0.3
 jupyter-sphinx>=0.3.2
 myst-nb
 tensorflow-datasets
+tensorflow
 dm-haiku
 flax
 jupytext
+scikit-learn