Update to new PT Theme

stack-info: PR: #2361, branch: drisspg/stack/76

Author: drisspg

docs/requirements.txt

@@ -1,8 +1,14 @@
-sphinx-gallery>0.11
-sphinx==5.0.0
-sphinx_design
-sphinx_copybutton
-sphinx-tabs
-matplotlib
-myst-parser
--e git+https://github.com/pytorch/pytorch_sphinx_theme.git#egg=pytorch_sphinx_theme
+sphinx==5.3.0
+-e git+https://github.com/pytorch/pytorch_sphinx_theme.git@pytorch_sphinx_theme2#egg=pytorch_sphinx_theme2
+sphinxcontrib.katex==0.8.6
+sphinxext-opengraph==0.9.1
+breathe==4.34.0  # only if generating C++
+exhale==0.2.3  # only if generating C++ docs
+docutils==0.17.1  # Changed from 0.16 to match sphinx-tabs requirement
+sphinx-design==0.4.0
+sphinxcontrib-mermaid==1.0.0
+myst-parser==0.18.1  # if want to contribute in markdown
+sphinx-gallery==0.11.1  # if hosting interactive tutorials
+sphinx-tabs==3.4.0
+sphinx-copybutton==0.5.2
+sphinx-sitemap==2.6.0

docs/source/_static/css/custom.css

@@ -185,3 +185,25 @@ article.pytorch-article .sphx-glr-thumbnails .sphx-glr-thumbcontainer {
 article.pytorch-article div.section div.wy-table-responsive tbody td {
     width: 50%;
 }
+
+/* Mermaid diagram styling */
+.mermaid {
+    text-align: center;
+    margin: 1.5em 0;
+}
+
+/* Ensure mermaid diagrams are visible */
+div.mermaid {
+    display: block !important;
+    width: 100% !important;
+    min-height: 200px;
+    background: #f8f9fa;
+    border: 1px solid #dee2e6;
+    border-radius: 0.25rem;
+    padding: 1rem;
+}
+
+/* Fix for mermaid in pytorch theme */
+.pytorch-article .mermaid {
+    overflow: visible !important;
+}

docs/source/_templates/mermaid_init.html

@@ -0,0 +1,19 @@
+<!-- Load Mermaid JS -->
+<script src="https://cdn.jsdelivr.net/npm/mermaid@10/dist/mermaid.min.js"></script>
+<script>
+document.addEventListener('DOMContentLoaded', function() {
+    mermaid.initialize({
+        startOnLoad: true,
+        theme: 'default',
+        flowchart: {
+            useMaxWidth: true,
+            htmlLabels: true
+        }
+    });
+
+    // Force re-render of mermaid diagrams
+    setTimeout(function() {
+        mermaid.init();
+    }, 500);
+});
+</script>

docs/source/api_reference.rst

@@ -0,0 +1,12 @@
+API Reference
+=============
+
+.. toctree::
+   :maxdepth: 2
+
+   api_ref_intro
+   api_ref_dtypes
+   api_ref_quantization
+   api_ref_sparsity
+   api_ref_float8
+   api_ref_kernel

docs/source/conf.py

@@ -23,10 +23,11 @@
 import os
 import sys
 
-import pytorch_sphinx_theme
-from docutils.parsers import rst
+sys.path.insert(0, os.path.abspath("."))
+import pytorch_sphinx_theme2
 
-sys.path.append(os.path.abspath("."))
+html_theme = "pytorch_sphinx_theme2"
+html_theme_path = [pytorch_sphinx_theme2.get_html_theme_path()]
 
 # -- General configuration ------------------------------------------------
 
@@ -46,11 +47,16 @@
     "sphinx.ext.napoleon",
     "sphinx.ext.viewcode",
     "sphinx.ext.duration",
+    "sphinx.ext.linkcode",
     "sphinx_tabs.tabs",
     "sphinx_design",
     "sphinx_gallery.gen_gallery",
     "sphinx_copybutton",
     "myst_parser",
+    "sphinx_sitemap",
+    "sphinxcontrib.mermaid",
+    "pytorch_sphinx_theme2",
+    "sphinxext.opengraph",
 ]
 
 sphinx_gallery_conf = {
@@ -72,6 +78,10 @@
 napoleon_google_docstring = True
 project = "torchao"
 
+# -- OpenGraph Protocol settings --
+ogp_site_url = "http://pytorch.org/ao"
+ogp_image = "https://pytorch.org/assets/images/social-share.jpg"
+
 # Get TORCHAO_VERSION_DOCS during the build.
 torchao_version_docs = os.environ.get("TORCHAO_VERSION_DOCS", None)
 print(f"torchao_version_docs: {torchao_version_docs}")
@@ -91,8 +101,18 @@
 print(f"Version: {version}")
 html_title = " ".join((project, version, "documentation"))
 
+# Determine if this is a release build
+RELEASE = version != "main"
+
+# Configure version for switcher if you have multiple versions
+switcher_version = "main" if not RELEASE else version
+
 # Add any paths that contain templates here, relative to this directory.
-templates_path = ["_templates"]
+theme_variables = pytorch_sphinx_theme2.get_theme_variables()
+templates_path = [
+    "_templates",
+    os.path.join(os.path.dirname(pytorch_sphinx_theme2.__file__), "templates"),
+]
 
 # The suffix(es) of source filenames.
 # You can specify multiple suffix as a list of string:
@@ -102,6 +122,27 @@
     ".md": "markdown",
 }
 
+# Myst-markdown configuration
+myst_enable_extensions = [
+    "colon_fence",
+    "deflist",
+    "html_image",
+]
+
+# Mermaid configuration
+mermaid_version = "10.6.1"
+mermaid_output_format = "raw"
+mermaid_init_js = """
+mermaid.initialize({
+    startOnLoad: true,
+    theme: 'default',
+    flowchart: {
+        useMaxWidth: true,
+        htmlLabels: true
+    }
+});
+"""
+
 # The master toctree document.
 master_doc = "index"
 
@@ -133,25 +174,90 @@
 # The theme to use for HTML and HTML Help pages.  See the documentation for
 # a list of builtin themes.
 #
-html_theme = "pytorch_sphinx_theme"
-html_theme_path = [pytorch_sphinx_theme.get_html_theme_path()]
+# Theme configuration is set earlier in the file (lines 29-30)
 
 # Theme options are theme-specific and customize the look and feel of a theme
 # further.  For a list of options available for each theme, see the
 # documentation.
 #
 html_theme_options = {
-    "collapse_navigation": False,
+    "navigation_with_keys": False,
+    "analytics_id": "GTM-T8XT4PS",
+    "logo": {
+        "text": "",
+    },
+    "icon_links": [
+        {
+            "name": "X",
+            "url": "https://x.com/PyTorch",
+            "icon": "fa-brands fa-x-twitter",
+        },
+        {
+            "name": "GitHub",
+            "url": "https://github.com/pytorch/ao",
+            "icon": "fa-brands fa-github",
+        },
+        {
+            "name": "Discourse",
+            "url": "https://dev-discuss.pytorch.org/",
+            "icon": "fa-brands fa-discourse",
+        },
+        {
+            "name": "PyPi",
+            "url": "https://pypi.org/project/torchao/",
+            "icon": "fa-brands fa-python",
+        },
+    ],
+    "use_edit_page_button": True,
+    "navbar_center": "navbar-nav",
+    # Option 2: Display version in navbar (since torchao is relatively new)
+    "navbar_start": ["pytorch_version"],
     "display_version": True,
-    "logo_only": True,
-    "pytorch_project": "docs",
-    "navigation_with_keys": True,
+    # Uncomment below when you have multiple versions and a versions.json file
+    # "switcher": {
+    #     "json_url": "https://pytorch.org/ao/torchao-versions.json",
+    #     "version_match": switcher_version,
+    # },
+    # "show_version_warning_banner": True,
+    # Make the top banner less intrusive on scroll
+    "announcement": None,  # This removes any announcement banner
+    "navbar_sticky": "top",  # Makes navbar sticky but less intrusive
 }
 
 html_logo = "_static/img/pytorch-logo-dark.svg"
 
 html_css_files = ["css/custom.css"]
 
+# Add custom JavaScript for mermaid
+html_js_files = [
+    ("https://cdn.jsdelivr.net/npm/mermaid@10/dist/mermaid.min.js", {"async": "async"}),
+]
+
+# Base URL for sitemap generation
+html_baseurl = "https://pytorch.org/ao/"
+
+# Configure date info for "Created On | Last Updated" feature and theme variables
+html_context = {
+    "date_info": {
+        # Optional: Add paths to skip for performance optimization
+        "paths_to_skip": [
+            "gen_modules/",  # Skip auto-generated API reference modules
+            "tutorials/",  # Skip auto-generated tutorial gallery
+        ],
+    },
+    "theme_variables": theme_variables,
+    "display_github": True,
+    "github_url": "https://github.com",
+    "github_user": "pytorch",
+    "github_repo": "ao",
+    "feedback_url": "https://github.com/pytorch/ao",
+    "github_version": "main",
+    "doc_path": "docs/source",
+    "library_links": theme_variables.get("library_links", []),
+    "community_links": theme_variables.get("community_links", []),
+    "language_bindings_links": html_theme_options.get("language_bindings_links", []),
+}
+
 # Add any paths that contain custom static files (such as style sheets) here,
 # relative to this directory. They are copied after the builtin static files,
 # so a file named "default.css" will overwrite the builtin "default.css".
@@ -176,8 +282,52 @@
 # See http://stackoverflow.com/a/41184353/3343043
 
 
-from custom_directives import CustomCardEnd, CustomCardItem, CustomCardStart
+# Custom directives are now handled by pytorch_sphinx_theme2
+# No need to register them manually
+
+
+import inspect
+
+
+def linkcode_resolve(domain, info):
+    """Link API objects to GitHub source code."""
+    if domain != "py":
+        return None
+    if not info["module"]:
+        return None
+
+    try:
+        # Import torchao module
+        import torchao
+
+        module = __import__(info["module"], fromlist=[""])
+        obj = module
+        for part in info["fullname"].split("."):
+            obj = getattr(obj, part)
+        # Get the source file and line number
+        obj = inspect.unwrap(obj)
+        fn = inspect.getsourcefile(obj)
+        _, lineno = inspect.getsourcelines(obj)
+    except Exception:
+        return None
+
+    # Get the relative path from the torchao package
+    try:
+        fn = os.path.relpath(fn, start=os.path.dirname(torchao.__file__))
+    except Exception:
+        return None
+
+    # Determine the tag/branch based on the version
+    if RELEASE and version != "main":
+        # For release versions, use the version tag
+        tag = f"v{version}"
+    else:
+        # For development versions, use main branch
+        tag = "main"
+
+    return f"https://github.com/pytorch/ao/blob/{tag}/torchao/{fn}#L{lineno}"
+
 
-rst.directives.register_directive("customcardstart", CustomCardStart)
-rst.directives.register_directive("customcarditem", CustomCardItem)
-rst.directives.register_directive("customcardend", CustomCardEnd)
+def setup(app):
+    """Configure Sphinx app for pytorch_sphinx_theme2 features."""
+    app.config.add_last_updated = True

docs/source/developer_notes.rst

@@ -0,0 +1,11 @@
+Developer Notes
+===============
+
+.. toctree::
+   :maxdepth: 2
+
+   quantization
+   sparsity
+   dtypes
+   performant_kernels
+   contributor_guide

docs/source/index.rst

@@ -7,38 +7,29 @@ using native PyTorch. Please checkout torchao `README <https://github.com/pytorc
 for an overall introduction to the library and recent highlight and updates.
 
 .. toctree::
-   :glob:
    :maxdepth: 1
-   :caption: Getting Started
+   :hidden:
+   :caption: Get Started
 
    quick_start
 
 .. toctree::
-   :glob:
    :maxdepth: 1
-   :caption: Developer Notes
+   :hidden:
+   :caption: Tutorials
 
-   quantization
-   sparsity
-   contributor_guide
+   tutorials
 
 .. toctree::
-   :glob:
    :maxdepth: 1
-   :caption: API Reference
+   :hidden:
+   :caption: API
 
-   api_ref_dtypes
-   api_ref_quantization
-   api_ref_sparsity
-   api_ref_float8
+   api_reference
 
 .. toctree::
-   :glob:
    :maxdepth: 1
-   :caption: Tutorials
+   :hidden:
+   :caption: Developer
 
-   serialization
-   subclass_basic
-   subclass_advanced
-   pretraining
-   torchao_vllm_integration
+   developer_notes