feat(batch): x batch is now of shape (batch_size, *schema.shape) rather than a flat vector

tachyonicClock · tachyonicClock · commit 4a5333f5e8b4 · 2026-05-01T09:13:59.000+12:00
diff --git a/notebooks/10_ocl.ipynb b/notebooks/10_ocl.ipynb
diff --git a/src/capymoa/base/_batch.py b/src/capymoa/base/_batch.py
@@ -16,14 +16,14 @@ class Batch(ABC):
     def batch_train(self, x: torch.Tensor, y: torch.Tensor) -> None:
         """Train the model with a batch of instances.
 
-        :param x: A batch of feature vectors of shape ``(batch_size, num_features)``.
+        :param x: A batch of feature vectors of shape ``(batch_size, *schema.shape)``.
         :param y: A batch of target values, typically a vector of shape ``(batch_size,)``.
         """
 
     @abstractmethod
     def batch_predict(self, x: torch.Tensor) -> torch.Tensor:
         """Predict the target values for a batch of instances.
 
-        :param x: A batch of feature vectors of shape ``(batch_size, num_features)``.
+        :param x: A batch of feature vectors of shape ``(batch_size, *schema.shape)``.
         :return: Predicted target values, typically a vector of shape ``(batch_size,)``.
         """
diff --git a/src/capymoa/base/_classifier.py b/src/capymoa/base/_classifier.py
@@ -179,7 +179,7 @@ def batch_predict(self, x: torch.Tensor) -> torch.Tensor:
 
     def train(self, instance: LabeledInstance) -> None:
         """Calls :func:`batch_train` with a batch of size 1."""
-        x = torch.from_numpy(instance.x).view(1, -1)
+        x = torch.from_numpy(instance.x).view(1, *self.schema.shape)
         x = x.to(self.device, self.x_dtype)
         y = torch.scalar_tensor(
             instance.y_index, dtype=self.y_dtype, device=self.device
@@ -188,7 +188,7 @@ def train(self, instance: LabeledInstance) -> None:
 
     def predict_proba(self, instance: Instance) -> Optional[LabelProbabilities]:
         """Calls :func:`batch_predict_proba` with a batch of size 1."""
-        x = torch.from_numpy(instance.x.reshape(1, -1))
+        x = torch.from_numpy(instance.x.reshape(1, *self.schema.shape))
         x = x.to(self.device, self.x_dtype)
         return self.batch_predict_proba(x).flatten().numpy().astype(np.float64)
 
diff --git a/src/capymoa/classifier/_finetune.py b/src/capymoa/classifier/_finetune.py
@@ -65,7 +65,7 @@ def __init__(
             function that takes the model parameters and returns an optimizer.
         :param criterion: Loss function to use for training.
         :param device: Hardware for training.
-        :param random_seed: Seeds torch :py:func:`torch.manual_seed`.
+        :param random_seed: Seed for PyTorch.
         """
         super().__init__(schema, random_seed)
         # seed for reproducibility
diff --git a/src/capymoa/evaluation/evaluation.py b/src/capymoa/evaluation/evaluation.py
@@ -1040,8 +1040,10 @@ def prequential_evaluation(
         if isinstance(learner, Batch):
             # Collect a batch of instances and predict them all at once
             np_x = np.stack([instance.x for instance in batch])
-            torch_x = torch.from_numpy(np_x).to(
-                device=learner.device, dtype=learner.x_dtype
+            torch_x = (
+                torch.from_numpy(np_x)
+                .to(device=learner.device, dtype=learner.x_dtype)
+                .view(len(batch), *stream.get_schema().shape)
             )
             torch_y = torch.tensor(
                 yb_true, dtype=learner.y_dtype, device=learner.device
diff --git a/src/capymoa/ocl/evaluation.py b/src/capymoa/ocl/evaluation.py
@@ -328,11 +328,11 @@ def _abstain_prediction_uniform(rng: np.random.Generator, n_classes: int) -> Lab
 def _batch_test(rng: np.random.Generator, learner: Classifier, x: Tensor) -> np.ndarray:
     """Test a batch of instances using the learner."""
     batch_size = x.shape[0]
-    x = x.view(batch_size, -1)
     if isinstance(learner, BatchClassifier):
         x = x.to(dtype=learner.x_dtype, device=learner.device)
         return learner.batch_predict(x).cpu().numpy()
     else:
+        x = x.view(batch_size, -1)
         yb_pred = np.zeros(batch_size, dtype=int)
         for i in range(batch_size):
             instance = Instance.from_array(learner.schema, x[i].numpy())
@@ -345,16 +345,16 @@ def _batch_test(rng: np.random.Generator, learner: Classifier, x: Tensor) -> np.
         return yb_pred
 
 
-def _batch_train(learner: Classifier, x: Tensor, y: Tensor):
+def _batch_train(learner: Classifier, x: Tensor, y: Tensor, x_shape: Sequence[int]):
     """Train a batch of instances using the learner."""
-    batch_size = x.shape[0]
-    x = x.view(batch_size, -1)
+    bs = x.shape[0]
     if isinstance(learner, BatchClassifier):
-        x = x.to(dtype=learner.x_dtype, device=learner.device)
+        x = x.to(dtype=learner.x_dtype, device=learner.device).view(bs, *x_shape)
         y = y.to(dtype=learner.y_dtype, device=learner.device)
         learner.batch_train(x, y)
     else:
-        for i in range(batch_size):
+        x = x.view(bs, -1)
+        for i in range(bs):
             instance = LabeledInstance.from_array(
                 learner.schema, x[i].numpy(), int(y[i].item())
             )
@@ -441,7 +441,7 @@ def ocl_train_eval_loop(
                 yb: Tensor
                 pbar.update(1)
                 yb_pred = _batch_test(rng, learner, xb)
-                _batch_train(learner, xb, yb)
+                _batch_train(learner, xb, yb, learner.schema.shape)
 
                 for y, y_pred in zip(yb, yb_pred, strict=True):
                     online_eval.update(y.item(), y_pred)
diff --git a/src/capymoa/ocl/strategy/_der.py b/src/capymoa/ocl/strategy/_der.py
@@ -68,7 +68,7 @@ def __init__(
         self._buffer = ReservoirSampler(
             capacity=buffer_size,
             buffers=dict(
-                x=((schema.get_num_attributes(),), torch.float32),
+                x=(schema.shape, torch.float32),
                 z=((schema.get_num_classes(),), torch.float32),
                 y=((), torch.long),
             ),
@@ -82,14 +82,14 @@ def _train_step(self, x: Tensor, y: Tensor, update_buffer: bool) -> None:
         n = x.shape[0]
 
         # Update Buffer
-        x_t = self._augment(x.view(n, *self.schema.shape)).flatten(1)  # $x_t$
+        x_t = self._augment(x)  # $x_t$
         z = self._model(x_t)  # $z$
         if update_buffer:
             self._buffer.update(x=x, z=z, y=y)
 
         # Sample buffer and augment
         xp, zp, _ = self._buffer.sample(n).values()  # $x'$, $z'$, $y'$
-        x_tp = self._augment(xp.view(n, *self.schema.shape)).flatten(1)  # $x_t'$
+        x_tp = self._augment(xp)  # $x_t'$
 
         reg = self._alpha * self._logit_loss(zp, self._model(x_tp))  # Line 7
         loss = self._criterion(z, y) + reg  # Line 8
diff --git a/src/capymoa/ocl/strategy/_ewc.py b/src/capymoa/ocl/strategy/_ewc.py
@@ -179,9 +179,7 @@ def __init__(
         self._optimiser = optimiser
         self._model = model
         self._criterion = torch.nn.CrossEntropyLoss()
-        self._buffer = SlidingWindow.new_xybuffer(
-            fim_buffer, schema.get_num_attributes()
-        )
+        self._buffer = SlidingWindow.new_xybuffer(fim_buffer, schema.shape)
 
         # Buffers for anchoring the model
         self._anchor_params = BufferList(
diff --git a/src/capymoa/ocl/strategy/_experience_replay.py b/src/capymoa/ocl/strategy/_experience_replay.py
@@ -1,5 +1,4 @@
 import torch
-from collections import OrderedDict
 from torch import Tensor
 
 from capymoa.base import BatchClassifier
@@ -57,21 +56,16 @@ def __init__(
         super().__init__(learner.schema, learner.random_seed)
         #: The wrapped learner to be trained with experience replay.
         self.learner = learner
-        self._buffer = ReservoirSampler(
-            capacity=buffer_size,
-            buffers=OrderedDict(
-                {
-                    "x": ((self.schema.get_num_attributes(),), self.learner.x_dtype),
-                    "y": ((), self.learner.y_dtype),
-                }
-            ),
-            rng=torch.Generator().manual_seed(learner.random_seed),
+        self._buffer = ReservoirSampler.new_xybuffer(
+            buffer_size,
+            learner.schema.shape,
+            torch.Generator().manual_seed(self.random_seed),
         )
         self.repeat = repeat
 
     def batch_train(self, x: Tensor, y: Tensor) -> None:
         # update the buffer with the new data
-        self._buffer.update(x=x.detach(), y=y.detach())
+        self._buffer.update(x=x, y=y)
 
         for _ in range(self.repeat):
             # sample from the buffer and construct training batch
diff --git a/src/capymoa/ocl/strategy/_gdumb.py b/src/capymoa/ocl/strategy/_gdumb.py
@@ -43,7 +43,7 @@ def __init__(
 
         self.original_state_dict = model.state_dict()
         self.buffer = GreedySampler.new_xybuffer(
-            capacity, schema.get_num_attributes(), torch.Generator().manual_seed(seed)
+            capacity, schema.shape, torch.Generator().manual_seed(seed)
         )
         self.loss_func = nn.CrossEntropyLoss()
 
diff --git a/src/capymoa/ocl/strategy/_ncm.py b/src/capymoa/ocl/strategy/_ncm.py
@@ -78,7 +78,7 @@ def __init__(
 
     @torch.no_grad()
     def batch_train(self, x: Tensor, y: Tensor) -> None:
-        x = self._pre_processor(x)
+        x = self._pre_processor(x).flatten(start_dim=1)
 
         # Update mean and count
         for i in range(self.schema.get_num_classes()):
@@ -91,8 +91,7 @@ def batch_train(self, x: Tensor, y: Tensor) -> None:
 
     @torch.no_grad()
     def batch_predict_proba(self, x: Tensor) -> Tensor:
-        assert x.ndim == 2, "Input must be a 2D array (batch_size, features)"
-        x = self._pre_processor(x)
+        x = self._pre_processor(x).flatten(start_dim=1)
 
         # Calculate distances to class means
         distances = torch.cdist(x.unsqueeze(0), self._class_means.unsqueeze(0))
diff --git a/src/capymoa/ocl/strategy/_rar.py b/src/capymoa/ocl/strategy/_rar.py
@@ -92,13 +92,12 @@ def __init__(
         """
 
         super().__init__(learner.schema)
-        num_features = learner.schema.get_num_attributes()
         self.learner = learner
         self.augment = augment
         self.repeats = repeats
         self.coreset = ReservoirSampler.new_xybuffer(
             coreset_size,
-            num_features,
+            self.schema.shape,
             rng=torch.Generator().manual_seed(learner.random_seed),
         )
         self.shape = learner.schema.shape
diff --git a/src/capymoa/ocl/strategy/_slda.py b/src/capymoa/ocl/strategy/_slda.py
@@ -106,7 +106,7 @@ def __init__(
 
     @torch.no_grad()
     def batch_train(self, x: Tensor, y: Tensor) -> None:
-        x = self._pre_processor(x)
+        x = self._pre_processor(x).flatten(start_dim=1)
 
         # Update class means and counts
         for i in range(self.schema.get_num_classes()):
@@ -130,7 +130,7 @@ def batch_predict_proba(self, x: Tensor) -> Tensor:
                 device=self.device,
             )
 
-        x = self._pre_processor(x)
+        x = self._pre_processor(x).flatten(start_dim=1)
         covariance = self._covariance / self._count + self._ridge
         weights: Tensor = torch.linalg.solve(covariance, self._class_means.T).T
         bias = -0.5 * (self._class_means @ weights.T).diagonal()
diff --git a/src/capymoa/ocl/util/_replay.py b/src/capymoa/ocl/util/_replay.py
@@ -1,4 +1,4 @@
-from typing import Dict, OrderedDict, Tuple
+from typing import Dict, Sequence, OrderedDict, Tuple
 from typing_extensions import override
 from torch import Tensor, nn
 from torch.utils.data import TensorDataset
@@ -56,23 +56,10 @@ def device(self) -> torch.device:
     def __init__(
         self,
         capacity: int,
-        buffers: Dict[str, Tuple[Tuple[int, ...], torch.dtype]],
+        buffers: Dict[str, Tuple[Sequence[int], torch.dtype]],
         rng: torch.Generator = torch.default_generator,
     ) -> None:
         super().__init__()
-        # Type check `buffers`
-        if not isinstance(buffers, Dict):
-            raise TypeError(f"Buffers must be an dict, got {type(buffers)}")
-        for key, (shape, dtype) in buffers.items():
-            if not isinstance(key, str):
-                raise TypeError(f"Buffer keys must be strings, got {type(key)}")
-            if not isinstance(shape, tuple) or not all(
-                isinstance(dim, int) for dim in shape
-            ):
-                raise TypeError(f"Buffer shapes must be tuples of ints, got {shape}")
-            if not isinstance(dtype, torch.dtype):
-                raise TypeError(f"Buffer dtypes must be torch.dtype, got {dtype}")
-
         self._capacity = capacity
         self._rng = rng
         self._count = 0
@@ -88,12 +75,12 @@ def __init__(
     def new_xybuffer(
         cls,
         capacity: int,
-        features: int,
+        x_shape: Sequence[int],
         rng: torch.Generator = torch.default_generator,
     ) -> "ReplayBuffer":
-        """Standard buffer with features ``x: (n, features) f32`` and labels ``y: (n,) i64``."""
+        """Standard buffer with features ``x: (n, *x_shape) f32`` and labels ``y: (n,) i64``."""
         return cls(
-            capacity, {"x": ((features,), torch.float32), "y": ((), torch.long)}, rng
+            capacity, {"x": (x_shape, torch.float32), "y": ((), torch.long)}, rng
         )
 
 
diff --git a/tests/ocl/test_der.py b/tests/ocl/test_der.py
diff --git a/tests/ocl/test_l2p.py b/tests/ocl/test_l2p.py
@@ -4,9 +4,13 @@
 can mock a dummy ViT model to test that the training loop does not error at least.
 """
 
+import numpy as np
+
 from capymoa.ocl.strategy.l2p import L2P, L2PViT
 from capymoa.ocl.datasets import TinySplitMNIST
 from capymoa.ocl.evaluation import ocl_train_eval_loop
+from capymoa.base import BatchClassifier
+from capymoa.instance import LabeledInstance
 from torch import Tensor, nn
 import torch
 
@@ -41,3 +45,35 @@ def test_l2p():
         scenario.train_loaders(32),
         scenario.test_loaders(32),
     )
+
+
+class ShapeRecordingBatchClassifier(BatchClassifier):
+    def __init__(self, schema):
+        super().__init__(schema)
+        self.train_shapes = []
+        self.predict_shapes = []
+
+    def batch_train(self, x: Tensor, y: Tensor) -> None:
+        self.train_shapes.append(tuple(x.shape))
+
+    def batch_predict_proba(self, x: Tensor) -> Tensor:
+        self.predict_shapes.append(tuple(x.shape))
+        return torch.zeros((x.shape[0], self.schema.get_num_classes()))
+
+
+def test_batch_classifier_single_instance_preserves_schema_shape():
+    scenario = TinySplitMNIST()
+    xb, yb = next(iter(scenario.train_loaders(2)[0]))
+    learner = ShapeRecordingBatchClassifier(scenario.schema)
+    instance = LabeledInstance.from_array(
+        scenario.schema, xb[0].view(-1).numpy(), int(yb[0])
+    )
+
+    learner.train(instance)
+    prediction = learner.predict_proba(instance)
+
+    expected_shape = (1, *scenario.schema.shape)
+    assert learner.train_shapes == [expected_shape]
+    assert learner.predict_shapes == [expected_shape]
+    assert isinstance(prediction, np.ndarray)
+    assert prediction.shape == (scenario.schema.get_num_classes(),)
diff --git a/tests/ocl/test_strategy.py b/tests/ocl/test_strategy.py
@@ -44,6 +44,7 @@ def pre_processor() -> nn.Module:
     """Create a pre-processor for the schema."""
     torch.manual_seed(0)
     return nn.Sequential(
+        nn.Flatten(),
         nn.Linear(256, 512),
         nn.ReLU(),
     )
diff --git a/tests/test_batch.py b/tests/test_batch.py

Original file line number	Diff line number	Diff line change
`@@ -43,7 +43,7 @@ def __init__(`
`43`	`43`
`44`	`44`	`self.original_state_dict = model.state_dict()`
`45`	`45`	`self.buffer = GreedySampler.new_xybuffer(`
`46`		`- capacity, schema.get_num_attributes(), torch.Generator().manual_seed(seed)`
	`46`	`+ capacity, schema.shape, torch.Generator().manual_seed(seed)`
`47`	`47`	`)`
`48`	`48`	`self.loss_func = nn.CrossEntropyLoss()`
`49`	`49`