Move mkdocs-autoref-plugin to new external tool

docs/macros.py

@@ -8,9 +8,27 @@
 from pathlib import Path
 from typing import Any
 
+import torch
+
 if typing.TYPE_CHECKING:
     from mkdocs_macros.plugin import MacrosPlugin
 
+import lightning
+
+try:
+    from mkdocs_autoref_plugin.autoref_plugin import default_reference_sources
+
+    default_reference_sources.extend(
+        [
+            lightning.Trainer,
+            lightning.LightningModule,
+            lightning.LightningDataModule,
+            torch.nn.Module,
+        ]
+    )
+except ImportError:
+    pass
+
 logger = logging.getLogger(__name__)
 
 

project/datamodules/__init__.py

@@ -1,6 +1,6 @@
 """Datamodules (datasets + preprocessing + dataloading)
 
-See the :ref:`lightning.LightningDataModule` class for more information.
+See the `lightning.LightningDataModule` class for more information.
 """
 
 from .image_classification import ImageClassificationDataModule

project/datamodules/image_classification/imagenet.py

@@ -46,10 +46,10 @@ class ImageNetDataModule(ImageClassificationDataModule):
     - Made this a subclass of VisionDataModule
 
     Notes:
+
     - train_dataloader uses the train split of imagenet2012 and puts away a portion of it for the validation split.
     - val_dataloader uses the part of the train split of imagenet2012  that was not used for training via
         `num_imgs_per_val_class`
-        - TODO: needs to pass split='val' to UnlabeledImagenet.
     - test_dataloader uses the validation split of imagenet2012 for testing.
         - TODO: need to pass num_imgs_per_class=-1 for test dataset and split="test".
     """

pyproject.toml

@@ -45,9 +45,6 @@ dependencies = [
 readme = "README.md"
 requires-python = ">= 3.10"
 
-[project.entry-points."mkdocs.plugins"]
-custom_autoref_plugin = "project.utils.autoref_plugin:CustomAutoRefPlugin"
-
 [dependency-groups]
 dev = [
     "mktestdocs>=0.2.2",
@@ -75,6 +72,7 @@ docs = [
     "mkdocs-video>=1.5.0",
     "mkdocs-section-index>=0.3.9",
     "mkdocs-macros-plugin>=1.0.5",
+    "mkdocs-autoref-plugin",
 ]
 gpu = ["jax[cuda12]>=0.4.31"]
 
@@ -118,6 +116,7 @@ managed = true
 
 [tool.uv.sources]
 remote-slurm-executor = { git = "https://github.com/lebrice/remote-slurm-executor", branch = "master" }
+mkdocs-autoref-plugin = { git = "https://github.com/lebrice/mkdocs-autoref-plugin" }
 
 # [tool.uv.tasks]
 # serve-docs = "mkdocs serve"