Rename the algorithm tests class (wip)

Signed-off-by: Fabrice Normandin <[email protected]>

docs/generate_reference_docs.py

@@ -3,25 +3,28 @@
 
 
 import textwrap
+from logging import getLogger as get_logger
 from pathlib import Path
 
 import mkdocs_gen_files
 import mkdocs_gen_files.nav
 
 from project.utils.env_vars import REPO_ROOTDIR
 
+logger = get_logger(__name__)
+
 module = "project"
 modules = [
     "project/main.py",
     "project/experiment.py",
 ]
-submodules = [
-    "project.algorithms",
-    "project.configs",
-    "project.datamodules",
-    "project.networks",
-    "project.utils",
-]
+# submodules = [
+#     "project.algorithms",
+#     "project.configs",
+#     "project.datamodules",
+#     "project.networks",
+#     "project.utils",
+# ]
 
 
 def _get_import_path(module_path: Path) -> str:
@@ -42,7 +45,7 @@ def main():
 def add_doc_for_module(module_path: Path, nav: mkdocs_gen_files.nav.Nav) -> None:
     """TODO."""
 
-    assert module_path.is_dir() and (module_path / "__init__.py").exists(), module_path
+    assert module_path.is_dir()  # and (module_path / "__init__.py").exists(), module_path
 
     children = list(
         p
@@ -52,7 +55,7 @@ def add_doc_for_module(module_path: Path, nav: mkdocs_gen_files.nav.Nav) -> None
     for child_module_path in children:
         child_module_import_path = _get_import_path(child_module_path)
         doc_file = child_module_path.relative_to(REPO_ROOTDIR).with_suffix(".md")
-        write_doc_file = f"reference/{doc_file}"
+        write_doc_file = "reference" / doc_file
 
         nav[tuple(child_module_import_path.split("."))] = f"{doc_file}"
 
@@ -71,11 +74,12 @@ def add_doc_for_module(module_path: Path, nav: mkdocs_gen_files.nav.Nav) -> None
         p
         for p in module_path.iterdir()
         if p.is_dir()
-        and (p / "__init__.py").exists()
+        and ((p / "__init__.py").exists() or len(list(p.glob("*.py"))) > 0)
         and not p.name.endswith("_test")
         and not p.name.startswith((".", "__"))
     )
     for submodule in submodules:
+        logger.info(f"Creating doc for {submodule}")
         add_doc_for_module(submodule, nav)
 
 

project/algorithms/testsuites/algorithm.py

@@ -1,13 +1,10 @@
 from typing import Literal, NotRequired, Protocol, TypedDict
 
 import torch
-from lightning import Callback, LightningModule, Trainer
+from lightning import Callback, Trainer
 from torch import Tensor
 from typing_extensions import TypeVar
 
-from project.datamodules.image_classification.image_classification import (
-    ImageClassificationDataModule,
-)
 from project.utils.types import PyTree
 from project.utils.types.protocols import DataModule, Module
 
@@ -39,9 +36,6 @@ class Algorithm(Module, Protocol[BatchType, StepOutputType]):
     datamodule: DataModule[BatchType]
     network: Module
 
-    example_input_array = LightningModule.example_input_array
-    _device: torch.device | None = None
-
     def __init__(
         self,
         *,
@@ -51,18 +45,13 @@ def __init__(
         super().__init__()
         self.datamodule = datamodule
         self.network = network
-        # fix for `self.device` property which defaults to cpu.
-        self._device = None
-
-        if isinstance(datamodule, ImageClassificationDataModule):
-            self.example_input_array = torch.zeros(
-                (datamodule.batch_size, *datamodule.dims), device=self.device
-            )
-
         self.trainer: Trainer
 
     def training_step(self, batch: BatchType, batch_index: int) -> StepOutputType:
-        """Performs a training step."""
+        """Performs a training step.
+
+        See `LightningModule.training_step` for more information.
+        """
         return self.shared_step(batch=batch, batch_index=batch_index, phase="train")
 
     def validation_step(self, batch: BatchType, batch_index: int) -> StepOutputType: