Switch to using named tuples to improve error handling (#169)

alan-cooney · web-flow · commit 5464e8fe9c57 · 2024-01-03T12:33:58.000-03:00
diff --git a/sparse_autoencoder/activation_resampler/activation_resampler.py b/sparse_autoencoder/activation_resampler/activation_resampler.py
@@ -1,5 +1,5 @@
 """Activation resampler."""
-from typing import Annotated
+from typing import Annotated, NamedTuple
 
 from einops import rearrange
 from jaxtyping import Bool, Float, Int64
@@ -22,6 +22,15 @@
 from sparse_autoencoder.train.utils import get_model_device
 
 
+class LossInputActivationsTuple(NamedTuple):
+    """Loss and corresponding input activations tuple."""
+
+    loss_per_item: Float[Tensor, Axis.names(Axis.BATCH, Axis.COMPONENT_OPTIONAL)]
+    input_activations: Float[
+        Tensor, Axis.names(Axis.BATCH, Axis.COMPONENT_OPTIONAL, Axis.INPUT_OUTPUT_FEATURE)
+    ]
+
+
 class ActivationResampler(AbstractActivationResampler):
     """Activation resampler.
 
@@ -182,10 +191,7 @@ def compute_loss_and_get_activations(
         autoencoder: SparseAutoencoder,
         loss_fn: AbstractLoss,
         train_batch_size: int,
-    ) -> tuple[
-        Float[Tensor, Axis.names(Axis.BATCH, Axis.COMPONENT_OPTIONAL)],
-        Float[Tensor, Axis.names(Axis.BATCH, Axis.COMPONENT_OPTIONAL, Axis.INPUT_OUTPUT_FEATURE)],
-    ]:
+    ) -> LossInputActivationsTuple:
         """Compute the loss on a random subset of inputs.
 
         Motivation:
@@ -226,18 +232,18 @@ def compute_loss_and_get_activations(
                 if batch_idx >= n_batches_required:
                     break
 
-            loss_result = torch.cat(loss_batches).to(model_device)
+            loss_per_item = torch.cat(loss_batches).to(model_device)
             input_activations = torch.cat(input_activations_batches).to(model_device)
 
             # Check we generated enough data
-            if len(loss_result) < n_inputs:
+            if len(loss_per_item) < n_inputs:
                 error_message = (
                     f"Cannot get {n_inputs} items from the store, "
-                    f"as only {len(loss_result)} were available."
+                    f"as only {len(loss_per_item)} were available."
                 )
                 raise ValueError(error_message)
 
-            return loss_result, input_activations
+            return LossInputActivationsTuple(loss_per_item, input_activations)
 
     @staticmethod
     def assign_sampling_probabilities(
@@ -440,7 +446,7 @@ def resample_dead_neurons(
 
             # Compute the loss for the current model on a random subset of inputs and get the
             # activations.
-            loss, input_activations = self.compute_loss_and_get_activations(
+            loss_per_item, input_activations = self.compute_loss_and_get_activations(
                 store=activation_store,
                 autoencoder=autoencoder,
                 loss_fn=loss_fn,
@@ -451,7 +457,7 @@ def resample_dead_neurons(
             # square of the autoencoder's loss on that input.
             sample_probabilities: Float[
                 Tensor, Axis.names(Axis.BATCH, Axis.COMPONENT_OPTIONAL)
-            ] = self.assign_sampling_probabilities(loss)
+            ] = self.assign_sampling_probabilities(loss_per_item)
 
             # For each dead neuron sample an input according to these probabilities.
             sampled_input: list[
diff --git a/sparse_autoencoder/autoencoder/abstract_autoencoder.py b/sparse_autoencoder/autoencoder/abstract_autoencoder.py
@@ -1,16 +1,30 @@
 """Abstract Sparse Autoencoder Model."""
 from abc import ABC, abstractmethod
+from typing import NamedTuple
 
 from jaxtyping import Float
 from torch import Tensor
-from torch.nn import Module, Parameter
+from torch.nn import Module
 
 from sparse_autoencoder.autoencoder.components.abstract_decoder import AbstractDecoder
 from sparse_autoencoder.autoencoder.components.abstract_encoder import AbstractEncoder
 from sparse_autoencoder.autoencoder.components.abstract_outer_bias import AbstractOuterBias
+from sparse_autoencoder.autoencoder.types import ResetOptimizerParameterDetails
 from sparse_autoencoder.tensor_types import Axis
 
 
+class AutoencoderForwardPassResult(NamedTuple):
+    """Autoencoder Forward Pass Result."""
+
+    learned_activations: Float[
+        Tensor, Axis.names(Axis.BATCH, Axis.COMPONENT_OPTIONAL, Axis.LEARNT_FEATURE)
+    ]
+
+    decoded_activations: Float[
+        Tensor, Axis.names(Axis.BATCH, Axis.COMPONENT_OPTIONAL, Axis.INPUT_OUTPUT_FEATURE)
+    ]
+
+
 class AbstractAutoencoder(Module, ABC):
     """Abstract Sparse Autoencoder Model.
 
@@ -42,7 +56,7 @@ def post_decoder_bias(self) -> AbstractOuterBias:
         """Post-decoder bias."""
 
     @property
-    def reset_optimizer_parameter_details(self) -> list[tuple[Parameter, int]]:
+    def reset_optimizer_parameter_details(self) -> list[ResetOptimizerParameterDetails]:
         """Reset optimizer parameter details.
 
         Details of the parameters that should be reset in the optimizer, when resetting
@@ -63,10 +77,7 @@ def forward(
         x: Float[
             Tensor, Axis.names(Axis.BATCH, Axis.COMPONENT_OPTIONAL, Axis.INPUT_OUTPUT_FEATURE)
         ],
-    ) -> tuple[
-        Float[Tensor, Axis.names(Axis.BATCH, Axis.COMPONENT_OPTIONAL, Axis.LEARNT_FEATURE)],
-        Float[Tensor, Axis.names(Axis.BATCH, Axis.COMPONENT_OPTIONAL, Axis.INPUT_OUTPUT_FEATURE)],
-    ]:
+    ) -> AutoencoderForwardPassResult:
         """Forward Pass.
 
         Args:
diff --git a/sparse_autoencoder/autoencoder/components/abstract_decoder.py b/sparse_autoencoder/autoencoder/components/abstract_decoder.py
@@ -8,6 +8,7 @@
 from torch import Tensor
 from torch.nn import Module, Parameter
 
+from sparse_autoencoder.autoencoder.types import ResetOptimizerParameterDetails
 from sparse_autoencoder.tensor_types import Axis
 
 
@@ -60,7 +61,7 @@ def weight(
 
     @property
     @abstractmethod
-    def reset_optimizer_parameter_details(self) -> list[tuple[Parameter, int]]:
+    def reset_optimizer_parameter_details(self) -> list[ResetOptimizerParameterDetails]:
         """Reset optimizer parameter details.
 
         Details of the parameters that should be reset in the optimizer, when resetting
diff --git a/sparse_autoencoder/autoencoder/components/abstract_encoder.py b/sparse_autoencoder/autoencoder/components/abstract_encoder.py
@@ -8,6 +8,7 @@
 from torch import Tensor
 from torch.nn import Module, Parameter
 
+from sparse_autoencoder.autoencoder.types import ResetOptimizerParameterDetails
 from sparse_autoencoder.tensor_types import Axis
 
 
@@ -66,7 +67,7 @@ def bias(self) -> Float[Parameter, Axis.names(Axis.COMPONENT_OPTIONAL, Axis.LEAR
 
     @property
     @abstractmethod
-    def reset_optimizer_parameter_details(self) -> list[tuple[Parameter, int]]:
+    def reset_optimizer_parameter_details(self) -> list[ResetOptimizerParameterDetails]:
         """Reset optimizer parameter details.
 
         Details of the parameters that should be reset in the optimizer, when resetting
diff --git a/sparse_autoencoder/autoencoder/components/linear_encoder.py b/sparse_autoencoder/autoencoder/components/linear_encoder.py
@@ -10,6 +10,7 @@
 from torch.nn import Parameter, ReLU, init
 
 from sparse_autoencoder.autoencoder.components.abstract_encoder import AbstractEncoder
+from sparse_autoencoder.autoencoder.types import ResetOptimizerParameterDetails
 from sparse_autoencoder.tensor_types import Axis
 from sparse_autoencoder.utils.tensor_shape import shape_with_optional_dimensions
 
@@ -63,7 +64,7 @@ def bias(self) -> Float[Parameter, Axis.names(Axis.COMPONENT_OPTIONAL, Axis.LEAR
         return self._bias
 
     @property
-    def reset_optimizer_parameter_details(self) -> list[tuple[Parameter, int]]:
+    def reset_optimizer_parameter_details(self) -> list[ResetOptimizerParameterDetails]:
         """Reset optimizer parameter details.
 
         Details of the parameters that should be reset in the optimizer, when resetting
@@ -73,7 +74,10 @@ def reset_optimizer_parameter_details(self) -> list[tuple[Parameter, int]]:
             List of tuples of the form `(parameter, axis)`, where `parameter` is the parameter to
             reset (e.g. encoder.weight), and `axis` is the axis of the parameter to reset.
         """
-        return [(self.weight, -2), (self.bias, -1)]
+        return [
+            ResetOptimizerParameterDetails(parameter=self.weight, axis=-2),
+            ResetOptimizerParameterDetails(parameter=self.bias, axis=-1),
+        ]
 
     activation_function: ReLU
     """Activation function."""
diff --git a/sparse_autoencoder/autoencoder/components/tests/test_abstract_decoder.py b/sparse_autoencoder/autoencoder/components/tests/test_abstract_decoder.py
@@ -9,6 +9,7 @@
 from torch.nn import Parameter, init
 
 from sparse_autoencoder.autoencoder.components.abstract_decoder import AbstractDecoder
+from sparse_autoencoder.autoencoder.types import ResetOptimizerParameterDetails
 from sparse_autoencoder.tensor_types import Axis
 
 
@@ -46,9 +47,9 @@ def weight(
         return self._weight
 
     @property
-    def reset_optimizer_parameter_details(self) -> list[tuple[Parameter, int]]:
+    def reset_optimizer_parameter_details(self) -> list[ResetOptimizerParameterDetails]:
         """Reset optimizer parameter details."""
-        return [(self.weight, 1)]
+        return [ResetOptimizerParameterDetails(parameter=self.weight, axis=1)]
 
     def forward(
         self, x: Float[Tensor, Axis.names(Axis.BATCH, Axis.COMPONENT, Axis.LEARNT_FEATURE)]
diff --git a/sparse_autoencoder/autoencoder/components/tests/test_abstract_encoder.py b/sparse_autoencoder/autoencoder/components/tests/test_abstract_encoder.py
@@ -11,6 +11,7 @@
 from sparse_autoencoder.autoencoder.components.abstract_encoder import (
     AbstractEncoder,
 )
+from sparse_autoencoder.autoencoder.types import ResetOptimizerParameterDetails
 from sparse_autoencoder.tensor_types import Axis
 
 
@@ -50,9 +51,12 @@ def reset_parameters(self) -> None:
         self._weight: Parameter = init.kaiming_normal_(self._weight)  # type: ignore
 
     @property
-    def reset_optimizer_parameter_details(self) -> list[tuple[Parameter, int]]:
+    def reset_optimizer_parameter_details(self) -> list[ResetOptimizerParameterDetails]:
         """Reset optimizer parameter details."""
-        return [(self.weight, 0), (self.bias, 0)]
+        return [
+            ResetOptimizerParameterDetails(parameter=self.weight, axis=0),
+            ResetOptimizerParameterDetails(parameter=self.bias, axis=0),
+        ]
 
 
 @pytest.fixture()
diff --git a/sparse_autoencoder/autoencoder/components/unit_norm_decoder.py b/sparse_autoencoder/autoencoder/components/unit_norm_decoder.py
@@ -9,6 +9,7 @@
 from torch.nn import Parameter, init
 
 from sparse_autoencoder.autoencoder.components.abstract_decoder import AbstractDecoder
+from sparse_autoencoder.autoencoder.types import ResetOptimizerParameterDetails
 from sparse_autoencoder.tensor_types import Axis
 from sparse_autoencoder.utils.tensor_shape import shape_with_optional_dimensions
 
@@ -65,7 +66,7 @@ def weight(
         return self._weight
 
     @property
-    def reset_optimizer_parameter_details(self) -> list[tuple[Parameter, int]]:
+    def reset_optimizer_parameter_details(self) -> list[ResetOptimizerParameterDetails]:
         """Reset optimizer parameter details.
 
         Details of the parameters that should be reset in the optimizer, when resetting
@@ -75,7 +76,7 @@ def reset_optimizer_parameter_details(self) -> list[tuple[Parameter, int]]:
             List of tuples of the form `(parameter, axis)`, where `parameter` is the parameter to
             reset (e.g. encoder.weight), and `axis` is the axis of the parameter to reset.
         """
-        return [(self.weight, -1)]
+        return [ResetOptimizerParameterDetails(parameter=self.weight, axis=-1)]
 
     @validate_call
     def __init__(
diff --git a/sparse_autoencoder/autoencoder/model.py b/sparse_autoencoder/autoencoder/model.py
@@ -8,7 +8,10 @@
 from torch import Tensor
 from torch.nn.parameter import Parameter
 
-from sparse_autoencoder.autoencoder.abstract_autoencoder import AbstractAutoencoder
+from sparse_autoencoder.autoencoder.abstract_autoencoder import (
+    AbstractAutoencoder,
+    AutoencoderForwardPassResult,
+)
 from sparse_autoencoder.autoencoder.components.linear_encoder import LinearEncoder
 from sparse_autoencoder.autoencoder.components.tied_bias import TiedBias, TiedBiasPosition
 from sparse_autoencoder.autoencoder.components.unit_norm_decoder import UnitNormDecoder
@@ -139,10 +142,7 @@ def forward(
         x: Float[
             Tensor, Axis.names(Axis.BATCH, Axis.COMPONENT_OPTIONAL, Axis.INPUT_OUTPUT_FEATURE)
         ],
-    ) -> tuple[
-        Float[Tensor, Axis.names(Axis.BATCH, Axis.COMPONENT_OPTIONAL, Axis.LEARNT_FEATURE)],
-        Float[Tensor, Axis.names(Axis.BATCH, Axis.COMPONENT_OPTIONAL, Axis.INPUT_OUTPUT_FEATURE)],
-    ]:
+    ) -> AutoencoderForwardPassResult:
         """Forward Pass.
 
         Args:
@@ -155,7 +155,8 @@ def forward(
         learned_activations = self._encoder(x)
         x = self._decoder(learned_activations)
         decoded_activations = self._post_decoder_bias(x)
-        return learned_activations, decoded_activations
+
+        return AutoencoderForwardPassResult(learned_activations, decoded_activations)
 
     def initialize_tied_parameters(self) -> None:
         """Initialize the tied parameters."""
diff --git a/sparse_autoencoder/autoencoder/types.py b/sparse_autoencoder/autoencoder/types.py
@@ -0,0 +1,18 @@
+"""Autoencoder types."""
+from typing import NamedTuple
+
+from torch.nn import Parameter
+
+
+class ResetOptimizerParameterDetails(NamedTuple):
+    """Reset Optimizer Parameter Details.
+
+    Details of a parameter that should be reset in the optimizer, when resetting
+    it's corresponding dictionary vectors.
+    """
+
+    parameter: Parameter
+    """Parameter to reset."""
+
+    axis: int
+    """Axis of the parameter to reset."""
diff --git a/sparse_autoencoder/loss/abstract_loss.py b/sparse_autoencoder/loss/abstract_loss.py
@@ -1,6 +1,6 @@
 """Abstract loss."""
 from abc import ABC, abstractmethod
-from typing import final
+from typing import NamedTuple, final
 
 from jaxtyping import Float
 from strenum import LowercaseStrEnum
@@ -22,6 +22,14 @@ class LossReductionType(LowercaseStrEnum):
     NONE = "none"
 
 
+class LossResultWithMetrics(NamedTuple):
+    """Loss result with any metrics to log."""
+
+    loss: Float[Tensor, Axis.COMPONENT] | Float[Tensor, Axis.SINGLE_ITEM]
+
+    loss_metrics: list[MetricResult]
+
+
 class AbstractLoss(Module, ABC):
     """Abstract loss interface.
 
@@ -119,7 +127,7 @@ def scalar_loss_with_log(
         ],
         batch_reduction: LossReductionType = LossReductionType.MEAN,
         component_reduction: LossReductionType = LossReductionType.NONE,
-    ) -> tuple[Float[Tensor, Axis.COMPONENT] | Float[Tensor, Axis.SINGLE_ITEM], list[MetricResult]]:
+    ) -> LossResultWithMetrics:
         """Scalar loss (reduced across the batch and component axis) with logging.
 
         Args:
@@ -181,7 +189,7 @@ def scalar_loss_with_log(
             case LossReductionType.NONE:
                 pass
 
-        return current_module_loss, metrics
+        return LossResultWithMetrics(loss=current_module_loss, loss_metrics=metrics)
 
     @final
     def __call__(
@@ -196,7 +204,7 @@ def __call__(
             Tensor, Axis.names(Axis.BATCH, Axis.COMPONENT_OPTIONAL, Axis.INPUT_OUTPUT_FEATURE)
         ],
         reduction: LossReductionType = LossReductionType.MEAN,
-    ) -> tuple[Float[Tensor, Axis.SINGLE_ITEM], list[MetricResult]]:
+    ) -> LossResultWithMetrics:
         """Batch scalar loss.
 
         Args:
diff --git a/sparse_autoencoder/loss/learned_activations_l1.py b/sparse_autoencoder/loss/learned_activations_l1.py
diff --git a/sparse_autoencoder/optimizer/abstract_optimizer.py b/sparse_autoencoder/optimizer/abstract_optimizer.py
diff --git a/sparse_autoencoder/train/pipeline.py b/sparse_autoencoder/train/pipeline.py
