Fixes following merge

Author: sgreenbury

src/auto_cast/decoders/base.py

@@ -8,14 +8,6 @@
 class Decoder(nn.Module, ABC):
     """Base Decoder."""
 
-    decoder_model: nn.Module
-    latent_dim: int
-
-    def __init__(self, latent_dim: int, output_channels: int) -> None:
-        super().__init__()
-        self.latent_dim = latent_dim
-        self.output_channels = output_channels
-
     def postprocess(self, decoded: Tensor) -> TensorBTSPlusC:
         """Optionally transform the decoded tensor before returning.
 

src/auto_cast/decoders/channels_last.py

@@ -1,7 +1,7 @@
 from einops import rearrange
 
 from auto_cast.decoders.base import Decoder
-from auto_cast.types import Tensor
+from auto_cast.types import Tensor, TensorBCWH, TensorBTWHC
 
 
 class ChannelsLast(Decoder):
@@ -33,3 +33,6 @@ def forward(self, x: Tensor) -> Tensor:
         )
         # Rearrange to channels-last: (B, C, T, W, H) -> (B, T, W, H, C)
         return rearrange(x, "b c t w h -> b t w h c")
+
+    def decode(self, z: TensorBCWH) -> TensorBTWHC:
+        return self.forward(z)

src/auto_cast/decoders/dc.py

@@ -86,8 +86,9 @@ def __init__(
         checkpointing: bool = False,
         identity_init: bool = True,
     ) -> None:
-        super().__init__(latent_dim=in_channels, output_channels=out_channels)
-
+        super().__init__()
+        self.latent_dim = in_channels
+        self.output_channels = out_channels
         attention_heads = attention_heads or {}
         assert len(hid_blocks) == len(hid_channels)
 

src/auto_cast/encoders/permute_concat.py

@@ -2,7 +2,7 @@
 from einops import rearrange
 
 from auto_cast.encoders.base import Encoder
-from auto_cast.types import Batch, Tensor
+from auto_cast.types import Batch, Tensor, TensorBCWH
 
 
 class PermuteConcat(Encoder):
@@ -27,3 +27,6 @@ def forward(self, batch: Batch) -> Tensor:
             scalars = scalars.expand(b, -1, t, w, h)
             x = torch.cat([x, scalars], dim=1)
         return rearrange(x, "b c t w h -> b (c t) w h")
+
+    def encode(self, batch: Batch) -> TensorBCWH:
+        return self.forward(batch)

src/auto_cast/models/encoder_decoder.py

@@ -1,3 +1,5 @@
+from typing import Any, Self
+
 import lightning as L
 import torch
 from torch import nn
@@ -18,6 +20,16 @@ class EncoderDecoder(L.LightningModule):
     def __init__(self):
         super().__init__()
 
+    @classmethod
+    def from_encoder_decoder(
+        cls, encoder: Encoder, decoder: Decoder, loss_func: nn.Module, **kwargs: Any
+    ) -> Self:
+        instance = cls(**kwargs)
+        instance.encoder = encoder
+        instance.decoder = decoder
+        instance.loss_func = loss_func
+        return instance
+
     def forward(self, batch: Batch) -> TensorBTSPlusC:
         return self.decoder(self.encoder(batch))
 

src/auto_cast/models/encoder_processor_decoder.py

@@ -7,7 +7,7 @@
 from auto_cast.models.encoder_decoder import EncoderDecoder
 from auto_cast.processors.base import Processor
 from auto_cast.processors.rollout import RolloutMixin
-from auto_cast.types import Batch, EncodedBatch, Tensor, TensorBTSPlusC
+from auto_cast.types import Batch, EncodedBatch, Tensor, TensorBMStarL, TensorBTSPlusC
 
 
 class EncoderProcessorDecoder(RolloutMixin[Batch], L.LightningModule):

src/auto_cast/types/__init__.py

@@ -39,6 +39,7 @@
 # Spatial only (no time dimension)
 TensorBCSPlus = Float[Tensor, "batch channel spatial *spatial"]
 TensorBWHC = Float[Tensor, "batch width height channel"]
+TensorBCWH = Float[Tensor, "batch channel width height"]
 TensorBWHDC = Float[Tensor, "batch width height depth channel"]
 TensorBSPlusC = Float[Tensor, "batch spatial *spatial channel"]
 

tests/decoders/test_channels_last.py

@@ -11,18 +11,13 @@ def test_channels_last_reorders_dimensions():
 
     # Input shape: (B, C*T, W, H) - simulating encoder output
     x = torch.randn(batch_size, channels * time_steps, width, height)
-
     output = decoder(x)
 
     # Expected output shape: (B, T, W, H, C)
     assert output.shape == (batch_size, time_steps, width, height, channels)
 
     # Verify the transformation is correct by checking a round-trip
-    # Create a known input in (B, T, W, H, C) format
     original = torch.randn(batch_size, time_steps, width, height, channels)
-    # Simulate encoder: (B, T, W, H, C) -> (B, C, T, W, H) -> (B, C*T, W, H)
     encoded = rearrange(original, "b t w h c -> b (c t) w h")
-    # Decode back
     decoded = decoder(encoded)
-    # Should match original
-    assert torch.allclose(decoded, original)
+    assert torch.allclose(decoded, original), "Decoded does not match original input"

tests/models/test_encoder_processor_decoder.py

@@ -35,7 +35,9 @@ def test_encoder_processor_decoder_training_step_runs(make_toy_batch, dummy_load
     encoder = PermuteConcat(with_constants=False)
     decoder = ChannelsLast(output_channels=output_channels, time_steps=time_steps)
     loss = nn.MSELoss()
-    encoder_decoder = EncoderDecoder(encoder=encoder, decoder=decoder, loss_func=loss)
+    encoder_decoder = EncoderDecoder.from_encoder_decoder(
+        encoder=encoder, decoder=decoder, loss_func=loss
+    )
 
     processor = TinyProcessor(in_channels=merged_channels)
     model = EncoderProcessorDecoder.from_encoder_processor_decoder(
@@ -67,7 +69,9 @@ def test_encoder_processor_decoder_rollout_is_mixin_backed(make_toy_batch):
     encoder = PermuteConcat(with_constants=False)
     decoder = ChannelsLast(output_channels=output_channels, time_steps=time_steps)
     loss = nn.MSELoss()
-    encoder_decoder = EncoderDecoder(encoder=encoder, decoder=decoder, loss_func=loss)
+    encoder_decoder = EncoderDecoder.from_encoder_decoder(
+        encoder=encoder, decoder=decoder, loss_func=loss
+    )
     processor = TinyProcessor(in_channels=merged_channels)
     model = EncoderProcessorDecoder.from_encoder_processor_decoder(
         encoder_decoder=encoder_decoder,

tests/models/test_vae.py

@@ -69,7 +69,9 @@ class _FlatDecoder(Decoder):
     """Minimal decoder that reconstructs flat tensors for tests."""
 
     def __init__(self, latent_dim: int, output_dim: int) -> None:
-        super().__init__(latent_dim=latent_dim, output_channels=output_dim)
+        super().__init__()
+        self.latent_dim = latent_dim
+        self.output_dim = output_dim
         self.net = nn.Sequential(
             nn.Linear(latent_dim, 2 * latent_dim),
             nn.GELU(),
@@ -111,7 +113,9 @@ class _FlatteningDecoder(Decoder):
     """Decoder that maps flat latents back to spatial tensors."""
 
     def __init__(self, latent_dim: int, output_shape: tuple[int, ...]) -> None:
-        super().__init__(latent_dim=latent_dim, output_channels=output_shape[0])
+        super().__init__()
+        self.latent_dim = latent_dim
+        self.output_channels = output_shape[0]
         self.output_shape = output_shape
         out_features = math.prod(output_shape)
         self.net = nn.Sequential(