Fix pre-existing CI failures

.github/workflows/test-doc.yml

@@ -5,7 +5,7 @@ on:
     paths:
       - 'docs/**'
       - 'aion/**'
-      - '.github/workflows/docs-check.yml'
+      - '.github/workflows/test-doc.yml'
       - 'pyproject.toml'
 
 jobs:

aion/codecs/quantizers/__init__.py

@@ -236,9 +236,11 @@ def forward(
         self, z_e: torch.Tensor
     ) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]:
         """Performs a forward pass through the vector quantizer.
+
         Args:
             z_e: Tensor (B, C, ...)
                 The input tensor to be quantized.
+
         Returns:
             z_q: Tensor
                 The quantized tensor.
@@ -314,9 +316,11 @@ def forward(
         self, z_e: torch.Tensor
     ) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]:
         """Performs a forward pass through the vector quantizer.
+
         Args:
             z_e: Tensor (B, C, ...)
                 The input tensor to be quantized.
+
         Returns:
             z_q: Tensor
                 The quantized tensor.

aion/codecs/quantizers/scalar.py

@@ -115,9 +115,11 @@ def forward(
         self, z_e: torch.Tensor
     ) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]:
         """Performs a forward pass through the vector quantizer.
+
         Args:
             z_e: torch.Tensor (B)
                 The input tensor to be quantized.
+
         Returns:
             z_q: torch.Tensor (B)
                 The quantized tensor.
@@ -473,9 +475,11 @@ def forward(
         self, z_e: torch.Tensor
     ) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]:
         """Performs a forward pass through the vector quantizer.
+
         Args:
             z_e: torch.Tensor (B, C, ...)
                 The input tensor to be quantized.
+
         Returns:
             z_q: torch.Tensor
                 The quantized tensor.
@@ -532,9 +536,11 @@ def forward(
         self, z_es: Dict[str, torch.Tensor]
     ) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]:
         """Performs a forward pass through the vector quantizer.
+
         Args:
             z_es: Dict[str, torch.Tensor]
                 The input tensor to be quantized.
+
         Returns:
             z_qs: torch.Tensor
                 The quantized tensor.
@@ -644,9 +650,11 @@ def forward(
         self, z_e: torch.Tensor
     ) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]:
         """Performs a forward pass through the vector quantizer.
+
         Args:
             z_e: torch.Tensor (B, C, ...)
                 The input tensor to be quantized.
+
         Returns:
             z_q: torch.Tensor
                 The quantized tensor.

aion/fourm/generate.py

@@ -1503,15 +1503,21 @@ def generate(
         seed=None,
     ):
         """Generates a sequence of tokens from the input modalities.
-        :param mod_dict: Dictionary of modalities.
-        :param schedule: Schedule of modalities to use.
-            List of dictionaries containing {target_domain, scheme, num_tokens, temperature, cfg_scale, cfg_cond_domains}.
-        :param top_k: top_k > 0: Keep only top k tokens with highest probability (a.k.a. top-k filtering).
-        :param top_p: top_p > 0.0: Keep the top tokens with cumulative probability >= top_p (a.k.a. nucleus filtering).
-        :param text_tokenizer: Text tokenizer.
-        :param verbose: Whether to print progress.
-        :param seed: Random seed.
-        :return: Generated mod dict.
+
+        Args:
+            mod_dict: Dictionary of modalities.
+            schedule: Schedule of modalities to use. A list of dictionaries containing
+                ``target_domain``, ``scheme``, ``num_tokens``, ``temperature``,
+                ``cfg_scale``, and ``cfg_cond_domains``.
+            top_k: If greater than zero, keep only the top-k tokens.
+            top_p: If greater than zero, keep tokens with cumulative probability at
+                least this value (nucleus filtering).
+            text_tokenizer: Text tokenizer.
+            verbose: Whether to print progress.
+            seed: Random seed.
+
+        Returns:
+            Generated modality dictionary.
         """
 
         # Input embedding -> tokenizes the modalities - Many are placeholder for now
@@ -1616,15 +1622,21 @@ def generate_iter(
         seed=None,
     ):
         """Iterator that generates a sequence of tokens from the input modalities step by step.
-        :param mod_dict: Dictionary of modalities.
-        :param schedule: Schedule of modalities to use.
-            List of dictionaries containing {target_domain, scheme, num_tokens, temperature, cfg_scale, cfg_cond_domains}.
-        :param top_k: top_k > 0: Keep only top k tokens with highest probability (a.k.a. top-k filtering).
-        :param top_p: top_p > 0.0: Keep the top tokens with cumulative probability >= top_p (a.k.a. nucleus filtering).
-        :param text_tokenizer: Text tokenizer.
-        :param verbose: Whether to print progress.
-        :param seed: Random seed.
-        :return: Iterator of generated mod dict.
+
+        Args:
+            mod_dict: Dictionary of modalities.
+            schedule: Schedule of modalities to use. A list of dictionaries containing
+                ``target_domain``, ``scheme``, ``num_tokens``, ``temperature``,
+                ``cfg_scale``, and ``cfg_cond_domains``.
+            top_k: If greater than zero, keep only the top-k tokens.
+            top_p: If greater than zero, keep tokens with cumulative probability at
+                least this value (nucleus filtering).
+            text_tokenizer: Text tokenizer.
+            verbose: Whether to print progress.
+            seed: Random seed.
+
+        Yields:
+            Generated modality dictionaries after each schedule step.
         """
 
         # Input embedding -> tokenizes the modalities - Many are placeholder for now

docs/api.rst

@@ -29,24 +29,39 @@ Base Classes
 Image Modalities
 ~~~~~~~~~~~~~~~~
 
-.. automodule:: aion.modalities
-   :members: LegacySurveyImage, HSCImage
+.. autoclass:: aion.modalities.LegacySurveyImage
+   :members:
+   :undoc-members:
+   :show-inheritance:
+
+.. autoclass:: aion.modalities.HSCImage
+   :members:
    :undoc-members:
    :show-inheritance:
 
 Spectrum Modalities
 ~~~~~~~~~~~~~~~~~~~
 
-.. automodule:: aion.modalities
-   :members: DESISpectrum, SDSSSpectrum
+.. autoclass:: aion.modalities.DESISpectrum
+   :members:
+   :undoc-members:
+   :show-inheritance:
+
+.. autoclass:: aion.modalities.SDSSSpectrum
+   :members:
    :undoc-members:
    :show-inheritance:
 
 Catalog Modalities
 ~~~~~~~~~~~~~~~~~~
 
-.. automodule:: aion.modalities
-   :members: LegacySurveyCatalog, LegacySurveySegmentationMap
+.. autoclass:: aion.modalities.LegacySurveyCatalog
+   :members:
+   :undoc-members:
+   :show-inheritance:
+
+.. autoclass:: aion.modalities.LegacySurveySegmentationMap
+   :members:
    :undoc-members:
    :show-inheritance:
 
@@ -56,43 +71,196 @@ Scalar Modalities
 Legacy Survey Scalars
 ^^^^^^^^^^^^^^^^^^^^^^
 
-.. automodule:: aion.modalities
-   :members: LegacySurveyFluxG, LegacySurveyFluxR, LegacySurveyFluxI, LegacySurveyFluxZ, LegacySurveyFluxW1, LegacySurveyFluxW2, LegacySurveyFluxW3, LegacySurveyFluxW4, LegacySurveyShapeR, LegacySurveyShapeE1, LegacySurveyShapeE2, LegacySurveyEBV
+.. autoclass:: aion.modalities.LegacySurveyFluxG
+   :members:
+   :undoc-members:
+   :show-inheritance:
+
+.. autoclass:: aion.modalities.LegacySurveyFluxR
+   :members:
+   :undoc-members:
+   :show-inheritance:
+
+.. autoclass:: aion.modalities.LegacySurveyFluxI
+   :members:
+   :undoc-members:
+   :show-inheritance:
+
+.. autoclass:: aion.modalities.LegacySurveyFluxZ
+   :members:
+   :undoc-members:
+   :show-inheritance:
+
+.. autoclass:: aion.modalities.LegacySurveyFluxW1
+   :members:
+   :undoc-members:
+   :show-inheritance:
+
+.. autoclass:: aion.modalities.LegacySurveyFluxW2
+   :members:
+   :undoc-members:
+   :show-inheritance:
+
+.. autoclass:: aion.modalities.LegacySurveyFluxW3
+   :members:
+   :undoc-members:
+   :show-inheritance:
+
+.. autoclass:: aion.modalities.LegacySurveyFluxW4
+   :members:
+   :undoc-members:
+   :show-inheritance:
+
+.. autoclass:: aion.modalities.LegacySurveyShapeR
+   :members:
+   :undoc-members:
+   :show-inheritance:
+
+.. autoclass:: aion.modalities.LegacySurveyShapeE1
+   :members:
+   :undoc-members:
+   :show-inheritance:
+
+.. autoclass:: aion.modalities.LegacySurveyShapeE2
+   :members:
+   :undoc-members:
+   :show-inheritance:
+
+.. autoclass:: aion.modalities.LegacySurveyEBV
+   :members:
    :undoc-members:
    :show-inheritance:
 
 HSC Scalars
 ~~~~~~~~~~~
 
-.. automodule:: aion.modalities
-   :members: HSCAG, HSCAR, HSCAI, HSCAZ, HSCAY, HSCMagG, HSCMagR, HSCMagI, HSCMagZ, HSCMagY, HSCShape11, HSCShape22, HSCShape12
+.. autoclass:: aion.modalities.HSCAG
+   :members:
+   :undoc-members:
+   :show-inheritance:
+
+.. autoclass:: aion.modalities.HSCAR
+   :members:
+   :undoc-members:
+   :show-inheritance:
+
+.. autoclass:: aion.modalities.HSCAI
+   :members:
+   :undoc-members:
+   :show-inheritance:
+
+.. autoclass:: aion.modalities.HSCAZ
+   :members:
+   :undoc-members:
+   :show-inheritance:
+
+.. autoclass:: aion.modalities.HSCAY
+   :members:
+   :undoc-members:
+   :show-inheritance:
+
+.. autoclass:: aion.modalities.HSCMagG
+   :members:
+   :undoc-members:
+   :show-inheritance:
+
+.. autoclass:: aion.modalities.HSCMagR
+   :members:
+   :undoc-members:
+   :show-inheritance:
+
+.. autoclass:: aion.modalities.HSCMagI
+   :members:
+   :undoc-members:
+   :show-inheritance:
+
+.. autoclass:: aion.modalities.HSCMagZ
+   :members:
+   :undoc-members:
+   :show-inheritance:
+
+.. autoclass:: aion.modalities.HSCMagY
+   :members:
+   :undoc-members:
+   :show-inheritance:
+
+.. autoclass:: aion.modalities.HSCShape11
+   :members:
+   :undoc-members:
+   :show-inheritance:
+
+.. autoclass:: aion.modalities.HSCShape22
+   :members:
+   :undoc-members:
+   :show-inheritance:
+
+.. autoclass:: aion.modalities.HSCShape12
+   :members:
    :undoc-members:
    :show-inheritance:
 
 Gaia Scalars
 ~~~~~~~~~~~~
 
-.. automodule:: aion.modalities
-   :members: GaiaFluxG, GaiaFluxBp, GaiaFluxRp, GaiaParallax, GaiaXpBp, GaiaXpRp
+.. autoclass:: aion.modalities.GaiaFluxG
+   :members:
+   :undoc-members:
+   :show-inheritance:
+
+.. autoclass:: aion.modalities.GaiaFluxBp
+   :members:
+   :undoc-members:
+   :show-inheritance:
+
+.. autoclass:: aion.modalities.GaiaFluxRp
+   :members:
+   :undoc-members:
+   :show-inheritance:
+
+.. autoclass:: aion.modalities.GaiaParallax
+   :members:
+   :undoc-members:
+   :show-inheritance:
+
+.. autoclass:: aion.modalities.GaiaXpBp
+   :members:
+   :undoc-members:
+   :show-inheritance:
+
+.. autoclass:: aion.modalities.GaiaXpRp
+   :members:
    :undoc-members:
    :show-inheritance:
 
 Coordinate Scalars
 ~~~~~~~~~~~~~~~~~~
 
-.. automodule:: aion.modalities
-   :members: Ra, Dec, Z
+.. autoclass:: aion.modalities.Ra
+   :members:
    :undoc-members:
    :show-inheritance:
 
-Utility Types
-~~~~~~~~~~~~~
+.. autoclass:: aion.modalities.Dec
+   :members:
+   :undoc-members:
+   :show-inheritance:
 
-.. automodule:: aion.modalities
-   :members: ScalarModalities, ModalityType
+.. autoclass:: aion.modalities.Z
+   :members:
    :undoc-members:
    :show-inheritance:
 
+Utility Types
+~~~~~~~~~~~~~
+
+.. py:data:: ScalarModalities
+
+   Mapping from scalar modality names to their corresponding modality classes.
+
+.. py:data:: ModalityType
+
+   Union type covering all supported modality data structures.
+
 Codec System
 ------------