Merge pull request #3631 from AI-Hypercomputer:fix_docs_format

PiperOrigin-RevId: 897468064

Author: Google-ML-Automation

docs/conf.py

@@ -51,6 +51,7 @@
     "sphinx_design",
     "sphinx_copybutton",
     "sphinx.ext.napoleon",
+    "sphinx.ext.mathjax",
     # This needs to be before autodoc^
     "sphinx.ext.autodoc",
     "sphinx.ext.autosummary",