Add notebook testing and a test mode flag.

Author: emilyfertig

.github/workflows/notebooks.yaml

@@ -0,0 +1,49 @@
+name: Notebooks
+
+on:
+  # Trigger the workflow on push or pull request, but only on dev branch
+  push:
+    branches:
+      - revamp-2025
+  pull_request:
+    branches:
+      - revamp-2025
+
+permissions:
+  contents: read  # to fetch code
+
+concurrency:
+  group: ${{ github.workflow }}-${{ github.head_ref || github.ref }}
+  cancel-in-progress: true
+
+jobs:
+  documentation:
+    name: Set up Python ${{ matrix.python-version }}
+    runs-on: ubuntu-latest
+    timeout-minutes: 10
+    strategy:
+      matrix:
+        python-version: ["3.12"]
+      fail-fast: false
+
+    steps:
+    - uses: actions/checkout@08c6903cd8c0fde910a37f88322edcfb5dd907a8  # v5.0.0
+    - name: Set up Python ${{ matrix.python-version }}
+      uses: actions/setup-python@e797f83bcb11b83ae66e0230d6156d7c80228e7c  # v6.0.0
+      with:
+        python-version: ${{ matrix.python-version }}
+    - name: Install dependencies
+      run: |
+        python -m pip install --upgrade pip
+        pip install .[dev]
+        pip install ipykernel pytest-notebook
+        pip freeze
+    - name: Turn on test mode
+      run: |
+        find docs/source -type f \( -name "*.md" -o -name "*.ipynb" \) -exec sed -i "s/JAS_TEST_MODE = False/JAS_TEST_MODE = True/g" {} +
+    - name: Register Python3 kernel
+      run: |
+        python -m ipykernel install --user --name "python3" --display-name "Python 3"
+    - name: Test notebooks
+      run: |
+        pytest -n auto --tb=short --nb-force-regen --nb-test-files docs/source/ --ignore=docs/source/JAX_for_LLM_pretraining.ipynb --ignore=docs/source/data_loaders_on_cpu_with_jax.ipynb --ignore=docs/source/data_loaders_on_gpu_with_jax.ipynb --ignore=docs/source/JAX_for_PyTorch_users.ipynb --ignore=docs/source/JAX_porting_PyTorch_model.ipynb

docs/source/conf.py

@@ -67,6 +67,9 @@
     'JAX_transformer_text_classification.md',
     'data_loaders_on_cpu_with_jax.md',
     'data_loaders_on_gpu_with_jax.md',
+    # Ignored test files.
+    'test_mode_test.md',
+    'test_mode_test.ipynb',
 ]
 
 suppress_warnings = [
@@ -104,4 +107,5 @@
     'JAX_transformer_text_classification.ipynb',
     'data_loaders_on_cpu_with_jax.ipynb',
     'data_loaders_on_gpu_with_jax.ipynb',
+    'test_mode_test.ipynb',
 ]

docs/source/test_mode_test.ipynb

@@ -0,0 +1,32 @@
+{
+ "cells": [
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "# Tests that CI substitutes False with True\n",
+    "JAS_TEST_MODE = False\n",
+    "JAS_TEST_MODE"
+   ]
+  }
+ ],
+ "metadata": {
+  "jupytext": {
+   "formats": "ipynb,md",
+   "main_language": "python"
+  },
+  "language_info": {
+   "name": "python"
+  },
+  "nbsphinx": {
+   "orphan": true
+  },
+  "nbreg": {
+    "diff_ignore": ["/cells/0/execution_count", "/cells/0/outputs", "/metadata"]
+  }
+ },
+ "nbformat": 4,
+ "nbformat_minor": 5
+}

docs/source/test_mode_test.md

@@ -0,0 +1,17 @@
+---
+jupyter:
+  jupytext:
+    formats: ipynb,md
+    main_language: python
+    text_representation:
+      extension: .md
+      format_name: markdown
+      format_version: '1.3'
+      jupytext_version: 1.15.2
+---
+
+```python
+# Tests that CI substitutes False with True
+JAS_TEST_MODE = False
+JAS_TEST_MODE
+```