Revise loss API with generic BatchT input

Author: sgreenbury

src/auto_cast/models/encoder_processor_decoder.py

@@ -165,8 +165,7 @@ def __init__(
 
     def training_step(self, batch: Batch, batch_idx: int) -> Tensor:  # noqa: ARG002
         encoded_batch = self.encoder_decoder.encoder.encode_batch(batch)
-        output = self.processor.map(encoded_batch.encoded_inputs)
-        loss = self.processor.loss(output, encoded_batch.encoded_output_fields)
+        loss = self.processor.loss(encoded_batch)
         self.log(
             "train_loss", loss, prog_bar=True, batch_size=batch.input_fields.shape[0]
         )

src/auto_cast/models/processor.py

@@ -45,16 +45,14 @@ def forward(self, x: TensorBMStarL) -> TensorBMStarL:
         return self.processor.map(x)
 
     def training_step(self, batch: EncodedBatch, batch_idx: int) -> Tensor:  # noqa: ARG002
-        output = self.processor.map(batch.encoded_inputs)
-        loss = self.processor.loss(output, batch.encoded_output_fields)
+        loss = self.processor.loss(batch)
         self.log(
             "train_loss", loss, prog_bar=True, batch_size=batch.encoded_inputs.shape[0]
         )
         return loss
 
     def validation_step(self, batch: EncodedBatch, batch_idx: int) -> Tensor:  # noqa: ARG002
-        output = self.processor.map(batch.encoded_inputs)
-        loss = self.processor.loss(output, batch.encoded_output_fields)
+        loss = self.processor.loss(batch)
         self.log(
             "val_loss", loss, prog_bar=True, batch_size=batch.encoded_inputs.shape[0]
         )

src/auto_cast/nn/fno.py

@@ -4,7 +4,7 @@
 from torch import nn
 
 from auto_cast.processors.base import Processor
-from auto_cast.types import Tensor
+from auto_cast.types import EncodedBatch, Tensor
 
 
 @runtime_checkable
@@ -13,7 +13,7 @@ class _HasGridCache(Protocol):
     _res: Any | None
 
 
-class FNOProcessor(Processor):
+class FNOProcessor(Processor[EncodedBatch]):
     """Fourier Neural Operator Module.
 
     A discrete processor that uses a Fourier Neural Operator (FNO) to learn
@@ -91,5 +91,6 @@ def _apply(self, fn, recurse: bool = True):
     def map(self, x: Tensor) -> Tensor:
         return self(x)
 
-    def loss(self, output: Tensor, target: Tensor) -> Tensor:
-        return self.loss_func(output, target)
+    def loss(self, batch: EncodedBatch) -> Tensor:
+        output = self.map(batch.encoded_inputs)
+        return self.loss_func(output, batch.encoded_output_fields)

src/auto_cast/processors/base.py

@@ -1,12 +1,12 @@
 from abc import ABC, abstractmethod
-from typing import Any
+from typing import Any, Generic
 
 from torch import nn
 
-from auto_cast.types import Tensor
+from auto_cast.types import BatchT, Tensor
 
 
-class Processor(ABC, nn.Module):
+class Processor(ABC, nn.Module, Generic[BatchT]):
     """Processor Base Class."""
 
     def __init__(
@@ -27,7 +27,7 @@ def __init__(
             setattr(self, key, value)
 
     @abstractmethod
-    def loss(self, output: Tensor, target: Tensor) -> Tensor:
+    def loss(self, batch: BatchT) -> Tensor:
         """Compute loss between output and target."""
 
     @abstractmethod

src/auto_cast/processors/rollout.py

@@ -1,13 +1,11 @@
 from __future__ import annotations
 
 from abc import ABC, abstractmethod
-from typing import Generic, TypeVar
+from typing import Generic
 
 import torch
 
-from auto_cast.types import RolloutOutput, Tensor
-
-BatchT = TypeVar("BatchT")
+from auto_cast.types import BatchT, RolloutOutput, Tensor
 
 
 class RolloutMixin(ABC, Generic[BatchT]):

src/auto_cast/types/__init__.py

@@ -1,5 +1,6 @@
 from collections.abc import Sequence
 from dataclasses import dataclass
+from typing import TypeVar
 
 import torch
 from jaxtyping import Float
@@ -56,6 +57,9 @@
 # Rollout output type
 RolloutOutput = tuple[Tensor, None] | tuple[Tensor, Tensor]
 
+# Generic batch type variable
+BatchT = TypeVar("BatchT")
+
 
 @dataclass
 class Sample:

tests/models/test_encoder_processor_decoder.py

@@ -6,10 +6,10 @@
 from auto_cast.models.encoder_decoder import EncoderDecoder
 from auto_cast.models.encoder_processor_decoder import EncoderProcessorDecoder
 from auto_cast.processors.base import Processor
-from auto_cast.types import Tensor
+from auto_cast.types import EncodedBatch, Tensor
 
 
-class TinyProcessor(Processor):
+class TinyProcessor(Processor[EncodedBatch]):
     def __init__(self, in_channels: int = 1) -> None:
         super().__init__()
         self.conv = nn.Conv2d(
@@ -25,8 +25,9 @@ def forward(self, x: Tensor) -> Tensor:
     def map(self, x: Tensor) -> Tensor:
         return self(x)
 
-    def loss(self, output: Tensor, target: Tensor) -> Tensor:
-        return self.loss_func(output, target)
+    def loss(self, batch: EncodedBatch) -> Tensor:
+        outputs = self(batch.encoded_inputs)
+        return self.loss_func(outputs, batch.encoded_output_fields)
 
 
 def test_encoder_processor_decoder_training_step_runs(make_toy_batch, dummy_loader):

tests/processors/test_processors.py

@@ -24,7 +24,7 @@ def _toy_encoded_batch(
     )
 
 
-class _IdentityProcessor(Processor):
+class _IdentityProcessor(Processor[EncodedBatch]):
     def __init__(self) -> None:
         super().__init__(
             loss_func=nn.MSELoss(),
@@ -33,8 +33,8 @@ def __init__(self) -> None:
     def map(self, x: Tensor) -> Tensor:
         return x
 
-    def loss(self, output: Tensor, target: Tensor) -> Tensor:
-        return self.loss_func(output, target)
+    def loss(self, batch: EncodedBatch) -> Tensor:
+        return self.loss_func(batch.encoded_inputs, batch.encoded_output_fields)
 
 
 def test_processor_rollout_handles_encoded_batches():