Add flexibility to lower level API

Author: sgreenbury

src/auto_cast/decoders/base.py

@@ -1,10 +1,9 @@
-from torch import nn
+from typing import Any
 
-from auto_cast.types import Tensor
+from torch import nn
 
 
 class Decoder(nn.Module):
     """Base Decoder."""
 
-    # Q: Should decoder handle all these input types
-    def forward(self, x: Tensor) -> Tensor: ...
+    def forward(self, *args: Any, **kwargs: Any) -> Any: ...

src/auto_cast/encoders/base.py

@@ -1,12 +1,12 @@
-from torch import nn
+from typing import Any
 
-from auto_cast.types import Tensor
+from torch import nn
 
 
 class Encoder(nn.Module):
     """Base encoder."""
 
-    def forward(self, x: Tensor) -> Tensor:
+    def forward(self, *args: Any, **kwargs: Any) -> Any:
         """Forward Pass through the Encoder."""
         msg = "To implement."
         raise NotImplementedError(msg)

src/auto_cast/models/encoder_decoder.py

@@ -1,10 +1,12 @@
+from typing import Any
+
 import lightning as L
 import torch
 from torch import nn
 
 from auto_cast.decoders import Decoder
 from auto_cast.encoders import Encoder
-from auto_cast.processors.base import Preprocessor, Processor
+from auto_cast.processors.base import Preprocessor
 from auto_cast.types import Batch, Tensor
 
 
@@ -19,16 +21,22 @@ class EncoderDecoder(L.LightningModule):
     def __init__(self):
         pass
 
-    def forward(self, x: Tensor) -> Tensor:
-        return self.decoder(self.encoder(x))
+    def forward(self, *args: Any, **kwargs: Any) -> Any:
+        return self.decoder(self.encoder(*args, **kwargs))
 
     def training_step(self, batch: Batch, batch_idx: int) -> Tensor:  # noqa: ARG002
         x = self.preprocessor(batch)
         output = self(x)
         loss = self.loss_func(output, batch["output_fields"])
         return loss  # noqa: RET504
 
-    def encode_only(self, x: Batch) -> Tensor:
+    def validation_step(self, batch: Batch, batch_idx: int) -> Tensor: ...
+
+    def test_step(self, batch: Batch, batch_idx: int) -> Tensor: ...
+
+    def predict_step(self, batch: Batch, batch_idx: int) -> Tensor: ...
+
+    def encode(self, x: Batch) -> Tensor:
         x = self.preprocessor(x)
         return self.encoder(x)
 
@@ -45,25 +53,7 @@ def forward(self, x: Tensor) -> Tensor:
         x = self.decoder(z)
         return x  # noqa: RET504
 
-    def reparametrize(self, mu, log_var):
+    def reparametrize(self, mu: Tensor, log_var: Tensor) -> Tensor:
         std = torch.exp(0.5 * log_var)
         eps = torch.randn_like(std)
         return mu + eps * std
-
-
-class EncoderProcessorDecoder(L.LightningModule):
-    """Encoder-Processor-Decoder Model."""
-
-    encoder_decoder: EncoderDecoder
-    processor: Processor
-
-    def __init__(self): ...
-
-    def forward(self, x: Tensor) -> Tensor:
-        return self.encoder_decoder.decoder(
-            self.processor(self.encoder_decoder.encoder(x))
-        )
-
-    def training_step(self, batch: Batch, batch_idx: int) -> Tensor: ...
-
-    def configure_optmizers(self): ...

src/auto_cast/models/encoder_processor_decoder.py

@@ -1,6 +1,9 @@
+from typing import Any
+
 import lightning as L
 
 from auto_cast.models.encoder_decoder import EncoderDecoder
+from auto_cast.preprocessor.base import Preprocessor
 from auto_cast.processors.base import Processor
 from auto_cast.types import Batch, Tensor
 
@@ -10,14 +13,19 @@ class EncoderProcessorDecoder(L.LightningModule):
 
     encoder_decoder: EncoderDecoder
     processor: Processor
+    preprocessor: Preprocessor
 
     def __init__(self): ...
 
-    def forward(self, x: Tensor) -> Tensor:
+    def forward(self, *args: Any, **kwargs: Any) -> Any:
         return self.encoder_decoder.decoder(
-            self.processor(self.encoder_decoder.encoder(x))
+            self.processor(self.encoder_decoder.encoder(*args, **kwargs))
         )
 
-    def training_step(self, batch: Batch, batch_idx: int) -> Tensor: ...
+    def training_step(self, batch: Batch, batch_idx: int) -> Tensor:  # noqa: ARG002
+        x = self.preprocessor(batch)
+        output = self(x)
+        loss = self.processor.loss_func(output, batch["output_fields"])
+        return loss  # noqa: RET504
 
     def configure_optmizers(self): ...

src/auto_cast/preprocessor/base.py

@@ -1,12 +1,17 @@
-from torch import nn
+from typing import Any
 
-from auto_cast.types import Batch, Tensor
+from auto_cast.types import Batch
 
 
-class Preprocessor(nn.Module):
-    """Base Preprocessor."""
+class Preprocessor:
+    """Base Preprocessor.
 
-    def forward(self, x: Batch) -> Tensor:
+    This is not trainable but can combine the elements of the batch into a form that
+    can be passed to the call/forward of the models.
+
+    """
+
+    def __call__(self, x: Batch) -> Any:
         """Forward Pass through the Preprocessor."""
         msg = "To implement."
         raise NotImplementedError(msg)

src/auto_cast/processors/base.py

@@ -1,42 +1,33 @@
 from abc import ABC, abstractmethod
+from typing import Any
 
 import lightning as L
 import torch
-from torch import nn  # noqa: F401
+from torch import nn
 
 from auto_cast.preprocessor import Preprocessor
 from auto_cast.types import Batch, RolloutOutput, Tensor
 
 
-class Processor(L.LightningModule, ABC):
+class Processor(L.LightningModule):
     """Processor Base Class."""
 
     teacher_forcing_ratio: float
     stride: int
     max_rollout_steps: int
     preprocessor: Preprocessor
+    loss_func: nn.Module
 
-    def __init__(self):
-        pass
-
-    # Option 1
-    def forward(self, x: Tensor) -> Tensor:
+    def forward(self, *args: Any, **kwargs: Any) -> Any:
         """Forward pass through the Processor."""
         msg = "To implement."
         raise NotImplementedError(msg)
 
     def training_step(self, batch: Batch, batch_idx: int) -> Tensor:  # noqa: ARG002
         x = self.preprocessor(batch)
-        return self(x)
-
-    # # Option 2
-    # def forward(self, x: Batch) -> Tensor:
-    #     """Forward pass through the Processor."""
-    #     msg = "To implement."
-    #     raise NotImplementedError(msg)
-
-    # def training_step(self, batch: Batch, batch_idx: int):
-    #     self(batch)
+        output = self(x)
+        loss = self.loss_func(output, batch["output_fields"])
+        return loss  # noqa: RET504
 
     def configure_optmizers(self):
         pass
@@ -48,9 +39,7 @@ def rollout(self, batch: Batch) -> RolloutOutput:
         for _time_step in range(0, self.max_rollout_steps, self.stride):
             x = self.preprocessor(batch)
             pred_outs.append(self(x))
-            gt_outs.append(
-                batch["output_fields"]
-            )  # Q: this assumes we have output fields
+            gt_outs.append(batch["output_fields"])  # This assumes we have output fields
         return torch.stack(pred_outs), torch.stack(gt_outs)