chore: import upstream snapshot with attribution

docs/conf.py

@@ -0,0 +1,740 @@
+# Licensed to the Apache Software Foundation (ASF) under one
+# or more contributor license agreements.  See the NOTICE file
+# distributed with this work for additional information
+# regarding copyright ownership.  The ASF licenses this file
+# to you under the Apache License, Version 2.0 (the
+# "License"); you may not use this file except in compliance
+# with the License.  You may obtain a copy of the License at
+#
+#   http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing,
+# software distributed under the License is distributed on an
+# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+# KIND, either express or implied.  See the License for the
+# specific language governing permissions and limitations
+# under the License.
+# ruff: noqa: E402, E501, F401
+
+#
+# documentation build configuration file, created by
+# sphinx-quickstart on Thu Jul 23 19:40:08 2015.
+#
+# This file is execfile()d with the current directory set to its
+# containing dir.
+#
+# Note that not all possible configuration values are present in this
+# autogenerated file.
+#
+# All configuration values have a default; values that are commented out
+# serve to show the default.
+import gc
+import inspect
+import os
+import re
+import sys
+from functools import partial
+from hashlib import md5
+from importlib import import_module
+from pathlib import Path
+from textwrap import dedent, indent
+from unittest.mock import patch
+
+# If extensions (or modules to document with autodoc) are in another directory,
+# 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.
+curr_path = Path(__file__).expanduser().absolute().parent
+if curr_path.name == "_staging":
+    # Can't use curr_path.parent, because sphinx_gallery requires a relative path.
+    tvm_path = Path(os.pardir, os.pardir)
+else:
+    tvm_path = Path(os.pardir)
+
+sys.path.insert(0, str(tvm_path.resolve() / "python"))
+sys.path.insert(0, str(tvm_path.resolve() / "docs"))
+
+# -- General configuration ------------------------------------------------
+
+# General information about the project.
+project = "tvm"
+author = "Apache Software Foundation"
+copyright = f"2020 - 2026, {author}"
+github_doc_root = "https://github.com/apache/tvm/tree/main/docs/"
+
+os.environ["TVM_BUILD_DOC"] = "1"
+
+
+# Version information.
+import tvm
+from tvm import te, testing, topi
+
+# The version is derived from the Git tag by setuptools_scm at build time and exposed
+# as tvm.__version__ (see [tool.setuptools_scm] in pyproject.toml).
+version = tvm.__version__
+release = version
+
+
+def monkey_patch(module_name, func_name):
+    """Helper function for monkey-patching library functions.
+
+    Used to modify a few sphinx-gallery behaviors to make the "Open in Colab" button work correctly.
+    Should be called as a decorator with arguments. Note this behaves differently from unittest's
+    @mock.patch, as our monkey_patch decorator should be placed on the new version of the function.
+    """
+    module = import_module(module_name)
+    original_func = getattr(module, func_name)
+
+    def decorator(function):
+        updated_func = partial(function, real_func=original_func)
+        setattr(module, func_name, updated_func)
+        return updated_func
+
+    return decorator
+
+
+CURRENT_FILE_CONF = None
+
+
+@monkey_patch("sphinx_gallery.py_source_parser", "split_code_and_text_blocks")
+def split_code_and_text_blocks(source_file, return_node, real_func):
+    """Monkey-patch split_code_and_text_blocks to expose sphinx-gallery's file-level config.
+
+    It's kinda gross, but we need access to file_conf to detect the requires_cuda flag.
+    """
+    global CURRENT_FILE_CONF
+    file_conf, blocks, node = real_func(source_file, return_node)
+    CURRENT_FILE_CONF = file_conf
+    return (file_conf, blocks, node)
+
+
+# This header replaces the default sphinx-gallery one in sphinx_gallery/gen_rst.py.
+# Colab button has been temporarily disabled due to prebuilt packages unavailability.
+COLAB_HTML_HEADER = """
+.. DO NOT EDIT. THIS FILE WAS AUTOMATICALLY GENERATED BY
+.. TVM'S MONKEY-PATCHED VERSION OF SPHINX-GALLERY. TO MAKE
+.. CHANGES, EDIT THE SOURCE PYTHON FILE:
+.. "{python_file}"
+
+.. only:: html
+
+    .. note::
+        :class: sphx-glr-download-link-note
+
+        You can click :ref:`here <sphx_glr_download_{ref_name}>` to run the Jupyter notebook locally.
+
+.. rst-class:: sphx-glr-example-title
+
+.. _sphx_glr_{ref_name}:
+
+"""
+
+# Google Colab allows opening .ipynb files on GitHub by appending a GitHub path to this base URL.
+COLAB_URL_BASE = "https://colab.research.google.com/github"
+
+# The GitHub path where the site is automatically deployed by tvm-bot.
+IPYTHON_GITHUB_BASE = "apache/tvm-site/blob/asf-site/docs/_downloads/"
+
+# The SVG image of the "Open in Colab" button.
+BUTTON = (
+    "https://raw.githubusercontent.com/tlc-pack/web-data/main/images/utilities/colab_button.svg"
+)
+
+
+@monkey_patch("sphinx_gallery.gen_rst", "save_rst_example")
+def save_rst_example(
+    example_rst, example_file, time_elapsed, memory_used, gallery_conf, language, real_func
+):
+    """Monkey-patch save_rst_example to customize the tutorial header.
+
+    Note: Colab button has been temporarily disabled. The colab_url and button_svg
+    are still generated but not used in the header template.
+    """
+
+    # The url is the md5 hash of the notebook path.
+    example_fname = os.path.relpath(example_file, gallery_conf["src_dir"])
+    ref_fname = example_fname.replace(os.path.sep, "_")
+    notebook_path = example_fname[:-2] + "ipynb"
+    digest = md5(notebook_path.encode()).hexdigest()
+
+    # Fixed documentation versions must link to different (earlier) .ipynb notebooks.
+    # Note: colab_url is generated but not currently used in the header template.
+    colab_url = f"{COLAB_URL_BASE}/{IPYTHON_GITHUB_BASE}"
+    if "dev" not in version:
+        colab_url += version + "/"
+    colab_url += digest + "/" + os.path.basename(notebook_path)
+
+    new_header = COLAB_HTML_HEADER.format(
+        python_file=example_fname, ref_name=ref_fname, colab_url=colab_url, button_svg=BUTTON
+    )
+    with patch("sphinx_gallery.gen_rst.EXAMPLE_HEADER", new_header):
+        real_func(
+            example_rst, example_file, time_elapsed, memory_used, gallery_conf, language=language
+        )
+
+
+INCLUDE_DIRECTIVE_RE = re.compile(r"^([ \t]*)\.\. include::\s*(.+)\n", flags=re.M)
+COMMENT_DIRECTIVE_RE = re.compile(r"^\.\.(?: .*)?\n(?:(?:  .*)?\n)*", flags=re.M)
+ADMONITION_DIRECTIVE_RE = re.compile(r"^\.\. admonition:: *(.*)\n((?:(?:  .*)?\n)*)\n", flags=re.M)
+
+
+@monkey_patch("sphinx_gallery.notebook", "rst2md")
+def rst2md(text, gallery_conf, target_dir, heading_levels, real_func):
+    """Monkey-patch rst2md to support comments and some include directives.
+
+    Currently, only include directives without any parameters are supported. Also, note that in
+    reStructuredText any unrecognized explicit markup block is treated as a comment (see
+    https://docutils.sourceforge.io/docs/ref/rst/restructuredtext.html#comments).
+
+    For callouts, we only replace generic "admonition" directives. All others should be replaced by
+    sphinx-gallery's rst2md. Note that the "alert" and "alert-info" tags are support in most IPython
+    notebooks, but they render kinda funky on Colab.
+    """
+
+    def load_include(match):
+        full_path = os.path.join(target_dir, match.group(2))
+        with open(full_path) as f:
+            lines = f.read()
+        indented = indent(lines, match.group(1)) + "\n"
+        return indented
+
+    text = re.sub(INCLUDE_DIRECTIVE_RE, load_include, text)
+
+    # Replace generic, titled admonitions with indented text. Other admonitions (e.g. .. note::)
+    # will be handled by sphinx-gallery's
+    def rewrite_generic_admonition(match):
+        title, text = match.groups()
+        stripped_text = dedent(text).strip()
+        return f'<div class="alert alert-info"><h4>{title}</h4><p>{stripped_text}</p></div>'
+
+    text = re.sub(ADMONITION_DIRECTIVE_RE, rewrite_generic_admonition, text)
+
+    # Call the real function, and then strip any remaining directives (i.e. comments)
+    text = real_func(text, gallery_conf, target_dir, heading_levels)
+    text = re.sub(COMMENT_DIRECTIVE_RE, "", text)
+    return text
+
+
+def install_request_hook(gallery_conf, fname):
+    testing.utils.install_request_hook(tvm_path.resolve() / "tests" / "python" / "request_hook.py")
+
+
+INSTALL_TVM_DEV = """\
+%%shell
+# Installs the latest dev build of TVM from PyPI. If you wish to build
+# from source, see https://tvm.apache.org/docs/install/from_source.html
+pip install apache-tvm --pre"""
+
+INSTALL_TVM_FIXED = f"""\
+%%shell
+# Installs TVM version {version} from PyPI. If you wish to build
+# from source, see https://tvm.apache.org/docs/install/from_source.html
+pip install apache-tvm=={version}"""
+
+INSTALL_TVM_CUDA_DEV = """\
+%%markdown
+> **Note:** This tutorial requires a CUDA-enabled build of TVM.
+> Pre-built CUDA wheels are not currently available on PyPI.
+> Please build TVM from source with CUDA enabled before running this notebook:
+> https://tvm.apache.org/docs/install/from_source.html"""
+
+INSTALL_TVM_CUDA_FIXED = f"""\
+%%markdown
+> **Note:** This tutorial requires a CUDA-enabled build of TVM (version {version}).
+> Pre-built CUDA wheels are not currently available on PyPI.
+> Please build TVM from source with CUDA enabled before running this notebook:
+> https://tvm.apache.org/docs/install/from_source.html"""
+
+
+@monkey_patch("sphinx_gallery.gen_rst", "jupyter_notebook")
+def jupyter_notebook(script_blocks, gallery_conf, target_dir, real_func):
+    """Monkey-patch sphinx-gallery to add a TVM import block to each IPython notebook.
+
+    If we had only one import block, we could skip the patching and just set first_notebook_cell.
+    However, how we import TVM depends on if we are using a fixed or dev version, and whether we
+    will use the GPU.
+
+    Tutorials requiring a CUDA-enabled build of TVM should use the flag:
+    # sphinx_gallery_requires_cuda = True
+    """
+
+    requires_cuda = CURRENT_FILE_CONF.get("requires_cuda", False)
+    fixed_version = "dev" not in version
+
+    if fixed_version and requires_cuda:
+        install_block = INSTALL_TVM_CUDA_FIXED
+    elif fixed_version and not requires_cuda:
+        install_block = INSTALL_TVM_FIXED
+    elif not fixed_version and requires_cuda:
+        install_block = INSTALL_TVM_CUDA_DEV
+    else:
+        install_block = INSTALL_TVM_DEV
+
+    new_conf = {**gallery_conf, "first_notebook_cell": install_block}
+    return real_func(script_blocks, new_conf, target_dir)
+
+
+# Add any Sphinx extension module names here, as strings. They can be
+# extensions coming with Sphinx (named 'sphinx.ext.*') or your custom ones
+extensions = [
+    "sphinx.ext.autodoc",
+    "sphinx.ext.autosummary",
+    "sphinx.ext.intersphinx",
+    "sphinx.ext.napoleon",
+    "sphinx.ext.mathjax",
+    "sphinx_gallery.gen_gallery",
+    "autodocsumm",
+]
+
+# Add any paths that contain templates here, relative to this directory.
+templates_path = ["_templates"]
+
+# The suffix(es) of source filenames.
+# You can specify multiple suffix as a list of string:
+# source_suffix = ['.rst', '.md']
+source_suffix = [".rst", ".md"]
+
+# The encoding of source files.
+# source_encoding = 'utf-8-sig'
+
+# generate autosummary even if no references
+autosummary_generate = True
+
+# The main toctree document.
+main_doc = "index"
+
+# The language for content autogenerated by Sphinx. Refer to documentation
+# for a list of supported languages.
+#
+# This is also used if you do content translation via gettext catalogs.
+# Usually you set "language" from the command line for these cases.
+language = None
+
+# There are two options for replacing |today|: either, you set today to some
+# non-false value, then it is used:
+# today = ''
+# Else, today_fmt is used as the format for a strftime call.
+# today_fmt = '%B %d, %Y'
+
+# List of patterns, relative to source directory, that match files and
+# directories to ignore when looking for source files.
+exclude_patterns = ["_build", "_staging"]
+
+# The TIRx API pages autodoc modules (tvm.tirx, tvm.backend.cuda) that re-export
+# common IR types (PrimExpr, Op, ...) from several modules, which makes a handful
+# of autodoc'd cross references ambiguous. Silence that specific, benign category.
+suppress_warnings = ["ref.python"]
+
+# The reST default role (used for this markup: `text`) to use for all
+# documents.
+# default_role = None
+
+# If true, '()' will be appended to :func: etc. cross-reference text.
+# add_function_parentheses = True
+
+# If true, the current module name will be prepended to all description
+# unit titles (such as .. function::).
+# add_module_names = True
+
+# If true, sectionauthor and moduleauthor directives will be shown in the
+# output. They are ignored by default.
+# show_authors = False
+
+# The name of the Pygments (syntax highlighting) style to use.
+pygments_style = "sphinx"
+
+# A list of ignored prefixes for module index sorting.
+# modindex_common_prefix = []
+
+# If true, keep warnings as "system message" paragraphs in the built documents.
+# keep_warnings = False
+
+# If true, `todo` and `todoList` produce output, else they produce nothing.
+todo_include_todos = False
+
+# -- Options for HTML output ----------------------------------------------
+
+html_theme = "sphinx_book_theme"
+
+html_title = "Apache TVM"
+
+# 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".
+html_static_path = ["_static"]
+
+html_logo = "_static/img/tvm-logo-small.png"
+
+html_favicon = "_static/img/tvm-logo-square.png"
+
+# The Apache trademark/copyright footer is rendered through sphinx-book-theme's
+# ``extra_footer`` hook (see footer_html below). This mirrors how the tvm-ffi docs
+# (3rdparty/tvm-ffi/docs/conf.py) preserve the ASF menu under the book theme.
+footer_dropdown = {
+    "name": "ASF",
+    "items": [
+        ("Apache Homepage", "https://apache.org/"),
+        ("License", "https://www.apache.org/licenses/"),
+        ("Sponsorship", "https://www.apache.org/foundation/sponsorship.html"),
+        ("Security", "https://tvm.apache.org/docs/reference/security.html"),
+        ("Thanks", "https://www.apache.org/foundation/thanks.html"),
+        ("Events", "https://www.apache.org/events/current-event"),
+    ],
+}
+
+footer_note = " ".join(
+    """
+Copyright © 2026 The Apache Software Foundation. Apache TVM, Apache, the Apache feather,
+and the Apache TVM project logo are either trademarks or registered trademarks of
+the Apache Software Foundation.""".split("\n")
+).strip()
+
+
+def footer_html() -> str:
+    """Build the extra footer: ASF dropdown and the Apache trademark note.
+
+    The copyright line is rendered natively by sphinx-book-theme (from the ``copyright``
+    config value), so it is intentionally not repeated here.
+    """
+    dropdown_items = ""
+    for item_name, item_url in footer_dropdown["items"]:
+        dropdown_items += f'<li><a class="dropdown-item" href="{item_url}" target="_blank" style="font-size: 0.9em;">{item_name}</a></li>\n'
+
+    return f"""
+  <div class="footer-container" style="margin: 5px 0; font-size: 0.9em; color: #6c757d; text-align: right;">
+      <div class="footer-line1" style="display: flex; justify-content: flex-end; align-items: center; gap: 0.9em; margin-bottom: 3px; flex-wrap: wrap;">
+          <div class="footer-dropdown">
+              <div class="dropdown">
+                  <button class="btn btn-link dropdown-toggle" type="button" id="footerDropdown" data-bs-toggle="dropdown"
+                  aria-expanded="false" style="font-size: 0.9em; color: #6c757d; text-decoration: none; padding: 0; border: none; background: none;">
+                      {footer_dropdown["name"]}
+                  </button>
+                  <ul class="dropdown-menu dropdown-menu-end" aria-labelledby="footerDropdown" style="font-size: 0.9em;">
+{dropdown_items}                  </ul>
+              </div>
+          </div>
+      </div>
+      <div class="footer-line2" style="font-size: 0.9em; color: #6c757d; text-align: right;">
+          {footer_note}
+      </div>
+  </div>
+  """
+
+
+html_theme_options = {
+    "repository_url": "https://github.com/apache/tvm",
+    "repository_branch": "main",
+    "path_to_docs": "docs/",
+    "use_repository_button": True,
+    "use_edit_page_button": True,
+    "use_source_button": True,
+    "use_issues_button": True,
+    "show_toc_level": 2,
+    "show_navbar_depth": 1,
+    "icon_links": [
+        {
+            "name": "TVM Homepage",
+            "url": "https://tvm.apache.org/",
+            "icon": "fa-solid fa-house",
+            "type": "fontawesome",
+        },
+        {
+            "name": "Community",
+            "url": "https://tvm.apache.org/community",
+            "icon": "fa-solid fa-users",
+            "type": "fontawesome",
+        },
+        {
+            "name": "Download",
+            "url": "https://tvm.apache.org/download",
+            "icon": "fa-solid fa-box-open",
+            "type": "fontawesome",
+        },
+    ],
+    "extra_footer": footer_html(),
+}
+
+
+# Output file base name for HTML help builder.
+htmlhelp_basename = project + "doc"
+
+# -- Options for LaTeX output ---------------------------------------------
+latex_elements = {}
+
+# Grouping the document tree into LaTeX files. List of tuples
+# (source start file, target name, title,
+#  author, documentclass [howto, manual, or own class]).
+latex_documents = [
+    (main_doc, f"{project}.tex", project, author, "manual"),
+]
+
+intersphinx_mapping = {
+    "python": (f"https://docs.python.org/{sys.version_info.major}", None),
+    # "numpy": ("https://numpy.org/doc/stable", None),
+    # "scipy": ("https://docs.scipy.org/doc/scipy", None),
+    # "matplotlib": ("https://matplotlib.org/", None),
+}
+
+from sphinx_gallery.sorting import ExplicitOrder
+
+examples_dirs = [
+    tvm_path.joinpath("docs", "get_started", "tutorials"),
+    tvm_path.joinpath("docs", "how_to", "tutorials"),
+    tvm_path.joinpath("docs", "deep_dive", "relax", "tutorials"),
+    tvm_path.joinpath("docs", "deep_dive", "tensor_ir", "tutorials"),
+]
+
+gallery_dirs = [
+    "get_started/tutorials/",
+    "how_to/tutorials/",
+    "deep_dive/relax/tutorials/",
+    "deep_dive/tensor_ir/tutorials/",
+]
+
+# Explicitly define the order within a subsection.
+# The listed files are sorted according to the list.
+# The unlisted files are sorted by filenames.
+# The unlisted files always appear after listed files.
+within_subsection_order = {}
+
+
+class WithinSubsectionOrder:
+    def __init__(self, src_dir):
+        self.src_dir = src_dir.split("/")[-1]
+
+    def __call__(self, filename):
+        # If the order is provided, use the provided order
+        if (
+            self.src_dir in within_subsection_order
+            and filename in within_subsection_order[self.src_dir]
+        ):
+            index = within_subsection_order[self.src_dir].index(filename)
+            assert index < 1e10
+            return f"\0{index:010d}"
+
+        # Otherwise, sort by filename
+        return filename
+
+
+# When running the tutorials on GPUs we are dependent on the Python garbage collector
+# collecting TVM packed function closures for any device memory to also be released. This
+# is not a good setup for machines with lots of CPU ram but constrained GPU ram, so force
+# a gc after each example.
+def force_gc(gallery_conf, fname):
+    gc.collect()
+
+
+filename_pattern_default = ".*"
+
+sphinx_gallery_conf = {
+    "backreferences_dir": "gen_modules/backreferences",
+    "doc_module": ("tvm", "numpy"),
+    "reference_url": {
+        "tvm": None,
+        # "matplotlib": "https://matplotlib.org/",
+        # "numpy": "https://numpy.org/doc/stable",
+    },
+    "examples_dirs": examples_dirs,
+    "within_subsection_order": WithinSubsectionOrder,
+    "gallery_dirs": gallery_dirs,
+    "filename_pattern": os.environ.get("TVM_TUTORIAL_EXEC_PATTERN", filename_pattern_default),
+    "download_all_examples": False,
+    "min_reported_time": 60,
+    "expected_failing_examples": [],
+    "reset_modules": ("matplotlib", "seaborn", force_gc, install_request_hook),
+    "promote_jupyter_magic": True,
+    # Drop the "Gallery generated by Sphinx-Gallery" signature line on generated pages.
+    "show_signature": False,
+}
+
+autodoc_default_options = {
+    "member-order": "bysource",
+}
+
+# Maps the original namespace to list of potential modules
+# that we can import alias from.
+tvm_alias_check_map = {
+    "tvm.te": ["tvm.tirx"],
+    "tvm.tirx": ["tvm.ir", "tvm.runtime"],
+}
+
+
+def update_alias_docstring(name, obj, lines):
+    """Update the docstring of alias functions.
+
+    This function checks if the obj is an alias of another documented object
+    in a different module.
+
+    If it is an alias, then it will append the alias information to the docstring.
+
+    Parameters
+    ----------
+    name : str
+        The full name of the object in the doc.
+
+    obj : object
+        The original object.
+
+    lines : list
+        The docstring lines, need to be modified inplace.
+    """
+    arr = name.rsplit(".", 1)
+    if len(arr) != 2:
+        return
+    target_mod, target_name = arr
+
+    if target_mod not in tvm_alias_check_map:
+        return
+    if not hasattr(obj, "__module__"):
+        return
+    obj_mod = obj.__module__
+
+    for amod in tvm_alias_check_map[target_mod]:
+        if not obj_mod.startswith(amod):
+            continue
+
+        if hasattr(sys.modules[amod], target_name):
+            obj_type = ":py:func" if callable(obj) else ":py:class"
+            lines.append(f".. rubric:: Alias of {obj_type}:`{amod}.{target_name}`")
+
+
+tvm_class_name_rewrite_map = {
+    "tvm.tirx": ["Var", "Call"],
+    "tvm.relax": ["Var", "Call", "StringImm"],
+    "tvm.relax.frontend.nn": ["Module"],
+}
+
+# When documenting modules under these prefixes, prefer types from the mapped module
+# to resolve ambiguous cross-references (e.g. Var exists in both tvm.tirx and tvm.relax).
+tvm_module_type_preference = {
+    "tvm.s_tir": "tvm.tirx",
+}
+
+
+def distinguish_class_name(name: str, lines: list[str]):
+    """Distinguish the docstring of type annotations.
+
+    In the whole TVM, there are many classes with the same name but in different modules,
+    e.g. ``tirx.Var``, ``relax.Var``. This function is used to distinguish them in the docstring,
+    by adding the module name as prefix.
+
+    To be specific, this function will check the current object name, and if it in the specific
+    module with specific name, it will add the module name as prefix to the class name to prevent
+    the confusion. Further, we only add the prefix to those standalone class name, but skip
+    the pattern of `xx.Var`, `Var.xx` and `xx.Var.xx`.
+
+    Parameters
+    ----------
+    name : str
+        The full name of the object in the doc.
+
+    lines : list
+        The docstring lines, need to be modified inplace.
+    """
+    remap = {}
+    for module_name in tvm_class_name_rewrite_map:
+        if name.startswith(module_name):
+            short_name = module_name[4:] if module_name.startswith("tvm.") else module_name
+            for class_name in tvm_class_name_rewrite_map[module_name]:
+                remap.update({class_name: f"{short_name}.{class_name}"})
+
+    for k, v in remap.items():
+        for i in range(len(lines)):
+            lines[i] = re.sub(rf"(?<!\.)\b{k}\b(?!\.)", v, lines[i])
+
+
+def process_docstring(app, what, name, obj, options, lines):
+    """Sphinx callback to process docstring"""
+    if callable(obj) or inspect.isclass(obj):
+        update_alias_docstring(name, obj, lines)
+        distinguish_class_name(name, lines)
+
+
+def strip_ipython_magic(app, docname, source):
+    """Prevents IPython magic commands from being rendered in HTML files.
+
+    TODO rework this function to remove IPython magic commands from include directives too.
+    """
+    for i in range(len(source)):
+        source[i] = re.sub(r"%%.*\n\s*", "", source[i])
+
+
+def _patch_python_domain_find_obj():
+    """Patch PythonDomain.find_obj to resolve ambiguous cross-references.
+
+    Sphinx's ``warn-missing-reference`` event is only fired for unresolved
+    references. Ambiguous short names such as ``StringImm`` already have
+    multiple matches at ``PythonDomain.find_obj`` time, so the disambiguation
+    needs to happen here instead.
+    """
+    from sphinx.domains.python import PythonDomain
+
+    if getattr(PythonDomain.find_obj, "_tvm_patched", False):
+        return
+
+    _original_find_obj = PythonDomain.find_obj
+
+    def _common_prefix_len(lhs: str, rhs: str) -> int:
+        count = 0
+        for lpart, rpart in zip(lhs.split("."), rhs.split(".")):
+            if lpart != rpart:
+                break
+            count += 1
+        return count
+
+    def _dedup_find_obj(self, env, modname, classname, name, objtype, searchmode=0):
+        matches = _original_find_obj(self, env, modname, classname, name, objtype, searchmode)
+        if len(matches) <= 1:
+            return matches
+
+        short_name = name.rsplit(".", 1)[-1]
+
+        # Prefer a single canonical (non-aliased) entry if Sphinx already found one.
+        canonical_matches = [match for match in matches if not match[1].aliased]
+        if len(canonical_matches) == 1:
+            return canonical_matches
+
+        # Use TVM's module context for the known short names we rewrite in docstrings.
+        if modname:
+            candidate_modules = sorted(
+                (
+                    module_name
+                    for module_name, class_names in tvm_class_name_rewrite_map.items()
+                    if short_name in class_names and modname.startswith(module_name)
+                ),
+                key=len,
+                reverse=True,
+            )
+            for module_name in candidate_modules:
+                target_name = f"{module_name}.{short_name}"
+                context_matches = [match for match in matches if match[0] == target_name]
+                if len(context_matches) == 1:
+                    return context_matches
+
+            # Fall back to the unique match that best shares the current module prefix.
+            match_scores = {match[0]: _common_prefix_len(modname, match[0]) for match in matches}
+            best_score = max(match_scores.values())
+            if best_score > 1:
+                best_matches = [match for match in matches if match_scores[match[0]] == best_score]
+                if len(best_matches) == 1:
+                    return best_matches
+
+            # Check module type preference for cross-module resolution
+            # (e.g. tvm.s_tir.analysis uses types from tvm.tirx).
+            for prefix, preferred_mod in tvm_module_type_preference.items():
+                if modname.startswith(prefix):
+                    preferred = [m for m in matches if m[0].startswith(preferred_mod + ".")]
+                    if len(preferred) >= 1:
+                        return preferred[:1]
+
+        return matches
+
+    _dedup_find_obj._tvm_patched = True
+    PythonDomain.find_obj = _dedup_find_obj
+
+
+_patch_python_domain_find_obj()
+
+
+def setup(app):
+    app.connect("source-read", strip_ipython_magic)
+    app.connect("autodoc-process-docstring", process_docstring)