Update to new PT Theme

docs/README.md

@@ -0,0 +1,65 @@
+# torchao Documentation
+
+This directory contains the source files for the torchao documentation.
+
+## Building the Documentation
+
+### Prerequisites
+
+Install the required dependencies:
+
+```bash
+pip install -r requirements.txt
+pip install sphinx-serve
+```
+
+### Build the Documentation
+
+To build the HTML documentation:
+
+```bash
+make html
+```
+
+The built documentation will be in the `build/html` directory.
+
+### Serve the Documentation Locally
+
+After building, you can serve the documentation locally using `sphinx-serve`:
+
+```bash
+sphinx-serve -b build
+```
+
+This will start a local server (typically at http://localhost:8000) where you can view the documentation with live reload capabilities.
+
+## Documentation Structure
+
+- `source/` - Source files for the documentation
+  - `conf.py` - Sphinx configuration file
+  - `index.rst` - Main documentation entry point
+  - `tutorials.rst` - Tutorials section index
+  - `_static/` - Static files (CSS, images, etc.)
+  - `_templates/` - Custom templates
+  - `tutorials_source/` - Executable tutorial Python files (for sphinx-gallery)
+  - `tutorials/` - Generated tutorial gallery (auto-generated, don't edit)
+- `build/` - Generated documentation output (created after building)
+
+## Tutorial Types
+
+This documentation has two types of tutorials:
+
+### 1. Static Tutorials
+Educational content, guides, and explanations that are written as `.rst` or `.md` files:
+- Located directly in `source/` (e.g., `serialization.rst`, `subclass_basic.rst`)
+- Referenced in `source/tutorials.rst`
+- These are traditional documentation pages
+
+### 2. Executable Tutorials
+Interactive code examples and demos that can be run:
+- Source files: `source/tutorials_source/*.py` (Python scripts with special formatting)
+- Generated output: `source/tutorials/` (auto-generated HTML gallery)
+- Built using sphinx-gallery extension
+- Each Python file becomes a downloadable notebook and HTML page
+
+When you run `make html`, sphinx-gallery automatically converts Python files in `tutorials_source/` into an interactive gallery in the `tutorials/` directory.

docs/requirements.txt

@@ -1,9 +1,13 @@
-sphinx-gallery>0.11
-sphinx==5.0.0
-sphinx_design
-sphinx_copybutton
-sphinx-tabs
-matplotlib
-myst-parser
+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
+docutils==0.17.1  # Changed from 0.16 to match sphinx-tabs requirement
+sphinx-design==0.4.0
 sphinxcontrib-mermaid==1.0.0
--e git+https://github.com/pytorch/pytorch_sphinx_theme.git#egg=pytorch_sphinx_theme
+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
+myst-parser

docs/source/conf.py

@@ -16,17 +16,15 @@
 # add these directories to sys.path here. If the directory is relative to the
 # documentation root, use os.path.abspath to make it absolute, like shown here.
 #
-# import os
-# import sys
-# sys.path.insert(0, os.path.abspath('.'))
 
 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,18 +44,22 @@
     "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",
     "sphinxcontrib.mermaid",
+    "sphinx_sitemap",
+    "pytorch_sphinx_theme2",
+    "sphinxext.opengraph",
 ]
 
 sphinx_gallery_conf = {
-    "examples_dirs": "tutorials_source",  # path to your sphinx-gallery examples
-    "gallery_dirs": "tutorials",  # path to where to save shpinx-gallery generated output
-    "filename_pattern": "./*.py",  # any .py file in docs/source/tutorials will be built by sphinx-gallery
+    "examples_dirs": "tutorials_source",  # path to your sphinx-gallery examples source
+    "gallery_dirs": "tutorials",  # path to where to save sphinx-gallery generated output
+    "filename_pattern": "./*.py",  # any .py file in tutorials_source will be built by sphinx-gallery
     "backreferences_dir": "gen_modules/backreferences",  # path to store the backreferences
     "remove_config_comments": True,
 }
@@ -73,6 +75,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}")
@@ -92,8 +98,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:
@@ -134,25 +150,69 @@
 # 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",
+    "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,
+    "announcement": None,
 }
 
 html_logo = "_static/img/pytorch-logo-dark.svg"
 
 html_css_files = ["css/custom.css"]
 
+# Base URL for sitemap generation
+html_baseurl = "https://pytorch.org/ao/"
+
+# Configure date info for "Created On | Last Updated" feature
+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_user": "pytorch",
+    "github_repo": "ao",
+    "github_version": "main",
+    "doc_path": "docs/source",
+}
+
 # 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".
@@ -177,8 +237,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