apache--tvm/docs/conf.py

# 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)