Change the store to require a component dimension (#168)

sparse_autoencoder/__init__.py

@@ -1,6 +1,5 @@
 """Sparse Autoencoder Library."""
 from sparse_autoencoder.activation_resampler.activation_resampler import ActivationResampler
-from sparse_autoencoder.activation_store.disk_store import DiskActivationStore
 from sparse_autoencoder.activation_store.tensor_store import TensorActivationStore
 from sparse_autoencoder.autoencoder.model import SparseAutoencoder
 from sparse_autoencoder.loss.abstract_loss import LossReductionType
@@ -53,7 +52,6 @@
     "CapacityMetric",
     "Controller",
     "ControllerType",
-    "DiskActivationStore",
     "Distribution",
     "Goal",
     "HyperbandStopping",

sparse_autoencoder/activation_store/base_store.py

@@ -66,15 +66,15 @@ class ActivationStore(
     def append(
         self,
         item: Float[Tensor, Axis.names(Axis.INPUT_OUTPUT_FEATURE)],
-        component_idx: int = 0,
+        component_idx: int,
     ) -> Future | None:
         """Add a Single Item to the Store."""
 
     @abstractmethod
     def extend(
         self,
         batch: Float[Tensor, Axis.names(Axis.BATCH, Axis.INPUT_OUTPUT_FEATURE)],
-        component_idx: int = 0,
+        component_idx: int,
     ) -> Future | None:
         """Add a Batch to the Store."""
 
@@ -98,8 +98,8 @@ def __len__(self) -> int:
 
     @abstractmethod
     def __getitem__(
-        self, index: int
-    ) -> Float[Tensor, Axis.names(Axis.COMPONENT_OPTIONAL, Axis.INPUT_OUTPUT_FEATURE)]:
+        self, index: tuple[int, ...] | slice | int
+    ) -> Float[Tensor, Axis.names(Axis.ANY)]:
         """Get an Item from the Store."""
 
     def shuffle(self) -> None:
@@ -124,7 +124,7 @@ def fill_with_test_data(
 
         Example:
             >>> from sparse_autoencoder.activation_store.tensor_store import TensorActivationStore
-            >>> store = TensorActivationStore(max_items=100, n_neurons=256)
+            >>> store = TensorActivationStore(max_items=100, n_neurons=256, n_components=1)
             >>> store.fill_with_test_data(batch_size=100)
             >>> len(store)
             100