diff --git a/docs/source/conf.py b/docs/source/conf.py
index 6f9a3c8b..76bedf62 100644
--- a/docs/source/conf.py
+++ b/docs/source/conf.py
@@ -5,16 +5,17 @@
 # full list see the documentation:
 # http://www.sphinx-doc.org/en/master/config
 import glob
-import inspect
 import os
 import re
-import sys
-from typing import List, Optional
+from typing import List
 
 import lightning_utilities
 import pt_lightning_sphinx_theme
 from lightning_utilities.docs import adjust_linked_external_docs, fetch_external_assets
 
+# This function is used to populate the (source) links in the API
+from lightning_utilities.docs.formatting import _linkcode_resolve
+
 # -- Path setup --------------------------------------------------------------
 
 # If extensions (or modules to document with autodoc) are in another directory,
@@ -296,41 +297,10 @@ def _package_list_from_file(file: str) -> List[str]:
 
 
 # Resolve function
-# This function is used to populate the (source) links in the API
 def linkcode_resolve(domain, info):
-    def find_source() -> Optional[str]:
-        # try to find the file and line number, based on code from numpy:
-        # https://github.com/numpy/numpy/blob/master/doc/source/conf.py#L286
-        obj = sys.modules[info["module"]]
-        for part in info["fullname"].split("."):
-            obj = getattr(obj, part)
-        fname = inspect.getsourcefile(obj)
-        # https://github.com/rtfd/readthedocs.org/issues/5735
-        if any(s in fname for s in ("readthedocs", "rtfd", "checkouts")):
-            # /home/docs/checkouts/readthedocs.org/user_builds/pytorch_lightning/checkouts/
-            #  devel/pytorch_lightning/utilities/cls_experiment.py#L26-L176
-            path_top = os.path.abspath(os.path.join("..", "..", ".."))
-            fname = os.path.relpath(fname, start=path_top)
-        else:
-            # Local build, imitate master
-            fname = "main/" + os.path.relpath(fname, start=os.path.abspath(".."))
-        source, lineno = inspect.getsourcelines(obj)
-        return fname, lineno, lineno + len(source) - 1
-
-    if domain != "py" or not info["module"]:
-        return None
-    try:
-        filename = "%s#L%d-L%d" % find_source()
-    except Exception:
-        filename = info["module"].replace(".", "/") + ".py"
-    # import subprocess
-    # tag = subprocess.Popen(['git', 'rev-parse', 'HEAD'], stdout=subprocess.PIPE,
-    #                        universal_newlines=True).communicate()[0][:-1]
-    branch = filename.split("/")[0]
-    # do mapping from latest tags to main
-    branch = {"latest": "main", "stable": "main"}.get(branch, branch)
-    filename = "/".join([branch] + filename.split("/")[1:])
-    return f"https://github.com/{github_user}/{github_repo}/blob/{filename}"
+    return _linkcode_resolve(
+        domain, info, github_user=github_user, github_repo=github_repo, main_branch="main", stable_branch="main"
+    )
 
 
 autosummary_generate = True
diff --git a/src/lightning_utilities/docs/formatting.py b/src/lightning_utilities/docs/formatting.py
index 85f33636..07547e29 100644
--- a/src/lightning_utilities/docs/formatting.py
+++ b/src/lightning_utilities/docs/formatting.py
@@ -33,7 +33,14 @@ def _transform_changelog(path_in: str, path_out: str) -> None:
         fp.writelines(chlog_lines)
 
 
-def _linkcode_resolve(domain: str, github_user: str, github_repo: str, info: dict) -> str:
+def _linkcode_resolve(
+    domain: str,
+    info: dict,
+    github_user: str,
+    github_repo: str,
+    main_branch: str = "master",
+    stable_branch: str = "release/stable",
+) -> str:
     def find_source() -> Tuple[str, int, int]:
         # try to find the file and line number, based on code from numpy:
         # https://github.com/numpy/numpy/blob/master/doc/source/conf.py#L286
@@ -49,7 +56,7 @@ def find_source() -> Tuple[str, int, int]:
             fname = str(os.path.relpath(fname, start=path_top))
         else:
             # Local build, imitate master
-            fname = f'master/{os.path.relpath(fname, start=os.path.abspath(".."))}'
+            fname = f'{main_branch}/{os.path.relpath(fname, start=os.path.abspath(".."))}'
         source, line_start = inspect.getsourcelines(obj)
         return fname, line_start, line_start + len(source) - 1
 
@@ -64,7 +71,7 @@ def find_source() -> Tuple[str, int, int]:
     #                        universal_newlines=True).communicate()[0][:-1]
     branch = filename.split("/")[0]
     # do mapping from latest tags to master
-    branch = {"latest": "master", "stable": "master"}.get(branch, branch)
+    branch = {"latest": main_branch, "stable": stable_branch}.get(branch, branch)
     filename = "/".join([branch] + filename.split("/")[1:])
     return f"https://github.com/{github_user}/{github_repo}/blob/{filename}"