Merge pull request google#5222 from samanklesaria:doctest_check

PiperOrigin-RevId: 866671138

Author: Flax Authors

pyproject.toml

@@ -80,6 +80,7 @@ docs = [
     "sphinx-book-theme",
     "Pygments>=2.6.1",
     "ipykernel",
+    "tqdm==4.67.1",
     "myst_nb",
     "nbstripout",
     "recommonmark",