New plan - nbval

Author: awf

docs/source/04-benchmark.ipynb

@@ -35,42 +35,28 @@
   {
    "cell_type": "code",
    "execution_count": 2,
-   "metadata": {
-    "nbreg": {
-     "diff_replace": [
-      [
-       "/outputs",
-       ": +[0-9.]+ nsec .*",
-       ": TIME "
-      ]
-     ]
-    },
-    "tags": [
-     "NOTE: NBREG"
-    ]
-   },
+   "metadata": {},
    "outputs": [
     {
-     "name": "stdout",
+     "name": "stderr",
      "output_type": "stream",
      "text": [
-      "GFloat scalar                  :  5843.22 nsec (25 runs at size 10000)\n",
-      "GFloat vectorized, numpy arrays:    45.12 nsec (25 runs at size 1000000)\n",
-      "GFloat vectorized, JAX JIT     :     2.50 nsec (500 runs at size 1000000)\n",
-      "ML_dtypes                      :     2.42 nsec (500 runs at size 1000000)\n"
+      "An NVIDIA GPU may be present on this machine, but a CUDA-enabled jaxlib is not installed. Falling back to cpu.\n"
      ]
     },
     {
-     "name": "stderr",
+     "name": "stdout",
      "output_type": "stream",
      "text": [
-      "An NVIDIA GPU may be present on this machine, but a CUDA-enabled jaxlib is not installed. Falling back to cpu.\n"
+      "GFloat scalar                  :  6306.18 nsec (25 runs at size 10000)\n",
+      "GFloat vectorized, numpy arrays:    52.52 nsec (25 runs at size 1000000)\n",
+      "GFloat vectorized, JAX JIT     :     3.04 nsec (500 runs at size 1000000)\n",
+      "ML_dtypes                      :     2.69 nsec (500 runs at size 1000000)\n"
      ]
     }
    ],
    "source": [
-    "# Note cell contains additional metadata to make pytest-notebook ignore timing diffs\n",
-    "#  \"nbreg\": { \"diff_replace\": [[\"/outputs\", \": +[0-9.]+ nsec\", \": TIME \"]] }\n",
+    "# NBVAL_IGNORE_OUTPUT\n",
     "\n",
     "N = 1_000_000\n",
     "a = np.random.rand(N)\n",

pyproject.toml

@@ -40,5 +40,4 @@ module = "mx.*"
 ignore_missing_imports = true
 
 [tool.pytest.ini_options]
-nb_test_files = true
-nb_diff_ignore = ["/metadata/language_info"]
+addopts = "--nbval"

requirements-dev.txt

@@ -1,6 +1,6 @@
 # Requirements for tests (see "requirements-tests.txt" also for direct dependencies)
 pytest
-pytest-notebook
+nbval
 ml_dtypes
 jaxlib
 jax