[Docs] Make links to source link to source (pytorch#141186)

Rewrite [SOURCE] links in the API docs to point to the source file in github repo.

Pull Request resolved: pytorch#141186
Approved by: https://github.com/malfet, https://github.com/msaroufim

Co-authored-by: Nikita Shulga <[email protected]>

docs/source/conf.py

@@ -11,6 +11,8 @@
 # All configuration values have a default; values that are commented out
 # serve to show the default.
 
+import inspect
+
 # 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.
@@ -62,6 +64,7 @@
     "sphinx_copybutton",
     "sphinx_panels",
     "myst_parser",
+    "sphinx.ext.linkcode",
 ]
 
 # build the templated autosummary files
@@ -3378,6 +3381,45 @@
     html_title = " ".join((project, version, "documentation"))
     release = version
 
+
+# Use the linkcode extension to override [SOURCE] links to point
+# to the repo. Use the torch_version variable defined above to
+# determine link
+def linkcode_resolve(domain, info):
+    if domain != "py":
+        return None
+    if not info["module"]:
+        return None
+
+    try:
+        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)
+        source, lineno = inspect.getsourcelines(obj)
+    except Exception:
+        return None
+
+    # Determine the tag based on the torch_version
+    if RELEASE:
+        version_parts = torch_version.split(
+            "."
+        )  # For release versions, format as "vX.Y.Z" for correct path in repo
+        patch_version = (
+            version_parts[2].split("+")[0].split("a")[0]
+        )  # assuming a0 always comes after release version in versions.txt
+        version_path = f"v{version_parts[0]}.{version_parts[1]}.{patch_version}"
+    else:
+        version_path = torch.version.git_version
+    fn = os.path.relpath(fn, start=os.path.dirname(torch.__file__))
+    return (
+        f"https://github.com/pytorch/pytorch/blob/{version_path}/torch/{fn}#L{lineno}"
+    )
+
+
 # The language for content autogenerated by Sphinx. Refer to documentation
 # for a list of supported languages.
 #