diff --git a/bluemath_tk/deeplearning/_base_model.py b/bluemath_tk/deeplearning/_base_model.py
index 7b406ee..921c4dc 100644
--- a/bluemath_tk/deeplearning/_base_model.py
+++ b/bluemath_tk/deeplearning/_base_model.py
@@ -1,6 +1,7 @@
 import copy
 import inspect
 from abc import abstractmethod
+from numbers import Real
 
 import numpy as np
 import torch
@@ -8,6 +9,7 @@
 from tqdm import tqdm
 
 from ..core.models import BlueMathModel
+from .metrics import _validate_eps
 from .metrics import evaluate_reconstruction as evaluate_reconstruction_metric
 from .metrics import reconstruction_error as reconstruction_error_metric
 
@@ -74,7 +76,6 @@ def _build_model(self, *args, **kwargs) -> nn.Module:
 
         pass
 
-
     def _get_reconstruction_target(self, X: np.ndarray) -> np.ndarray:
         """Return the default reconstruction target for ``X``."""
         return X
@@ -97,11 +98,7 @@ def _batch_slices(
             for start in range(0, n_samples, batch_size)
         ]
 
-        if (
-            avoid_singleton
-            and len(slices) > 1
-            and slices[-1][1] - slices[-1][0] == 1
-        ):
+        if avoid_singleton and len(slices) > 1 and slices[-1][1] - slices[-1][0] == 1:
             previous_start, previous_stop = slices[-2]
             final_stop = slices[-1][1]
             previous_size = previous_stop - previous_start
@@ -120,8 +117,7 @@ def _requires_non_singleton_training_batches(self) -> bool:
         if self.model is None:
             return False
         return any(
-            isinstance(module, nn.BatchNorm1d)
-            for module in self.model.modules()
+            isinstance(module, nn.BatchNorm1d) for module in self.model.modules()
         )
 
     def _validate_or_set_build_input_shape(self, input_shape: tuple) -> None:
@@ -139,23 +135,100 @@ def _validate_or_set_build_input_shape(self, input_shape: tuple) -> None:
             )
 
     @staticmethod
+    def _validate_learning_rate(learning_rate: float) -> float:
+        """Return a finite, non-negative real scalar learning rate."""
+        if (
+            not isinstance(learning_rate, Real)
+            or isinstance(learning_rate, (bool, np.bool_))
+            or not np.isfinite(float(learning_rate))
+            or learning_rate < 0
+        ):
+            raise ValueError(
+                "learning_rate must be a finite, non-negative real scalar."
+            )
+        return float(learning_rate)
+
+    @staticmethod
+    def _validate_finite_array(array: np.ndarray, name: str) -> None:
+        """Require a finite, real-valued NumPy array."""
+        if not isinstance(array, np.ndarray):
+            raise TypeError(f"{name} must be a NumPy array.")
+        if not np.issubdtype(array.dtype, np.number):
+            raise TypeError(f"{name} must contain numeric values.")
+        if np.issubdtype(array.dtype, np.complexfloating):
+            raise TypeError(f"{name} must contain real-valued data.")
+        if not np.isfinite(array).all():
+            raise ValueError(f"{name} must contain only finite values.")
+        float32_limit = np.finfo(np.float32).max
+        if np.any(array > float32_limit) or np.any(array < -float32_limit):
+            raise ValueError(f"{name} must remain finite when converted to float32.")
+
+    def _validate_target_shape(
+        self,
+        X: np.ndarray,
+        target: np.ndarray,
+    ) -> None:
+        """Require the target to match the reconstruction shape exactly."""
+        if not isinstance(target, np.ndarray):
+            raise TypeError("y must be a NumPy array.")
+        if tuple(target.shape) != tuple(X.shape):
+            raise ValueError(
+                f"Target shape {target.shape} is incompatible with "
+                f"reconstruction shape {X.shape}."
+            )
+
     def _validate_inference_inputs(
+        self,
         X: np.ndarray,
         batch_size: int,
         name: str = "X",
+        check_expected_shape: bool = False,
     ) -> None:
-        """Validate common prediction, encoding, and decoding inputs."""
+        """Validate prediction, encoding, decoding, and metric inputs."""
         if not isinstance(X, np.ndarray):
             raise TypeError(f"{name} must be a NumPy array.")
         if X.ndim < 1:
             raise ValueError(f"{name} must have at least one dimension.")
         if len(X) == 0:
             raise ValueError(f"{name} must contain at least one sample.")
-        if batch_size < 1:
-            raise ValueError("batch_size must be at least 1.")
+        if (
+            not isinstance(batch_size, int)
+            or isinstance(batch_size, bool)
+            or batch_size < 1
+        ):
+            raise ValueError("batch_size must be a positive integer.")
+        self._validate_finite_array(X, name)
+
+        if check_expected_shape and self._build_input_shape is not None:
+            expected = tuple(self._build_input_shape[1:])
+            actual = tuple(X.shape[1:])
+            if actual != expected:
+                raise ValueError(f"Expected per-sample shape {expected}, got {actual}.")
+
+    def _validate_latent_inputs(
+        self,
+        Z: np.ndarray,
+        batch_size: int,
+    ) -> None:
+        """Require latent data with the public ``(batch, k)`` shape."""
+        expected_width = getattr(self, "k", None)
+        if Z.ndim != 2:
+            expected = (
+                f"(batch, {expected_width})"
+                if expected_width is not None
+                else "(batch, latent_width)"
+            )
+            raise ValueError(f"Z must have shape {expected}; got shape {Z.shape}.")
+        if Z.shape[0] < 1:
+            raise ValueError("Z must contain at least one latent vector.")
+        if expected_width is not None and Z.shape[1] != expected_width:
+            raise ValueError(
+                f"Z must have latent width {expected_width}; got {Z.shape[1]}."
+            )
+        self._validate_inference_inputs(Z, batch_size, name="Z")
 
-    @staticmethod
     def _validate_fit_inputs(
+        self,
         X: np.ndarray,
         y: np.ndarray,
         validation_split: float,
@@ -173,19 +246,33 @@ def _validate_fit_inputs(
                 "X must include a leading sample dimension. "
                 "For tabular data use shape (n_samples, n_features)."
             )
+        if any(dimension < 1 for dimension in X.shape[1:]):
+            raise ValueError("Every per-sample dimension must be positive.")
         if len(X) != len(y):
             raise ValueError(
-                f"X and y must contain the same number of samples; "
+                "X and y must contain the same number of samples; "
                 f"got {len(X)} and {len(y)}."
             )
-        if not 0.0 < validation_split < 1.0:
-            raise ValueError("validation_split must be strictly between 0 and 1.")
-        if batch_size < 1:
-            raise ValueError("batch_size must be at least 1.")
-        if epochs < 1:
-            raise ValueError("epochs must be at least 1.")
-        if patience < 1:
-            raise ValueError("patience must be at least 1.")
+        self._validate_target_shape(X, y)
+        self._validate_finite_array(X, "X")
+        self._validate_finite_array(y, "y")
+
+        if (
+            not isinstance(validation_split, (int, float))
+            or isinstance(validation_split, bool)
+            or not np.isfinite(float(validation_split))
+            or not 0.0 < validation_split < 1.0
+        ):
+            raise ValueError(
+                "validation_split must be finite and strictly between 0 and 1."
+            )
+        for name, value in (
+            ("batch_size", batch_size),
+            ("epochs", epochs),
+            ("patience", patience),
+        ):
+            if not isinstance(value, int) or isinstance(value, bool) or value < 1:
+                raise ValueError(f"{name} must be a positive integer.")
 
         split = int((1 - validation_split) * len(X))
         if split < 2:
@@ -194,9 +281,7 @@ def _validate_fit_inputs(
                 "Increase the dataset size or reduce validation_split."
             )
         if len(X) - split < 1:
-            raise ValueError(
-                "The validation split must contain at least one sample."
-            )
+            raise ValueError("The validation split must contain at least one sample.")
 
     def _get_init_config(self) -> dict:
         """Collect constructor parameters needed to recreate this model."""
@@ -223,13 +308,222 @@ def _get_init_config(self) -> dict:
 
     @staticmethod
     def _require_scalar_loss(loss: torch.Tensor) -> None:
-        """Raise when a criterion returns a non-scalar training loss."""
+        """Raise when a criterion does not return one scalar tensor."""
+        if not isinstance(loss, torch.Tensor):
+            raise TypeError("The training criterion must return a PyTorch tensor.")
         if loss.ndim != 0:
             raise ValueError(
                 "The training criterion must return a scalar loss. "
                 "Use reduction='mean' or reduction='sum'."
             )
 
+    @staticmethod
+    def _require_finite_tensor(value: torch.Tensor, phase: str) -> None:
+        """Reject non-tensor or non-finite model results."""
+        if not isinstance(value, torch.Tensor):
+            raise TypeError(f"{phase} must be a PyTorch tensor.")
+        if not torch.isfinite(value).all():
+            raise FloatingPointError(f"{phase} is not finite.")
+
+    @classmethod
+    def _require_finite_loss(cls, loss: torch.Tensor, phase: str) -> None:
+        """Abort immediately when a training objective is not finite."""
+        cls._require_finite_tensor(loss, f"{phase} loss")
+
+    @staticmethod
+    def _require_matching_output_shape(
+        output: torch.Tensor,
+        target: torch.Tensor,
+        phase: str,
+    ) -> None:
+        """Reject non-tensor or broadcastable model outputs."""
+        if not isinstance(output, torch.Tensor):
+            raise TypeError("Model output must be a PyTorch tensor.")
+        if not isinstance(target, torch.Tensor):
+            raise TypeError("Training target must be a PyTorch tensor.")
+        if tuple(output.shape) != tuple(target.shape):
+            raise ValueError(
+                f"{phase} output shape {tuple(output.shape)} does not match "
+                f"target shape {tuple(target.shape)}."
+            )
+
+    def _require_finite_gradients(self) -> None:
+        """Abort when any model gradient contains NaN or infinity."""
+        if self.model is None:
+            raise ValueError("Model must be built before checking gradients.")
+        for name, parameter in self.model.named_parameters():
+            gradient = parameter.grad
+            if gradient is not None and not torch.isfinite(gradient).all():
+                raise FloatingPointError(
+                    f"Gradient for parameter {name!r} is not finite."
+                )
+
+    def _require_finite_buffers(self) -> None:
+        """Abort when a registered model buffer is not finite."""
+        if self.model is None:
+            raise ValueError("Model must be built before checking buffers.")
+        for name, buffer in self.model.named_buffers():
+            if not torch.isfinite(buffer).all():
+                raise FloatingPointError(f"Buffer {name!r} is not finite.")
+
+    def _require_finite_parameters(self) -> None:
+        """Abort when model parameters or registered buffers are not finite."""
+        if self.model is None:
+            raise ValueError("Model must be built before checking parameters.")
+        for name, parameter in self.model.named_parameters():
+            if not torch.isfinite(parameter).all():
+                raise FloatingPointError(f"Parameter {name!r} is not finite.")
+        self._require_finite_buffers()
+
+    @staticmethod
+    def _require_finite_state_dict(
+        state_dict: dict,
+        phase: str = "Checkpoint",
+    ) -> None:
+        """Reject non-finite floating or complex checkpoint tensors."""
+        if not isinstance(state_dict, dict):
+            raise TypeError("model_state_dict must be a dictionary.")
+        for name, value in state_dict.items():
+            if not isinstance(value, torch.Tensor):
+                raise TypeError(f"{phase} state entry {name!r} must be a tensor.")
+            if (value.is_floating_point() or value.is_complex()) and not (
+                torch.isfinite(value).all()
+            ):
+                raise FloatingPointError(f"{phase} state entry {name!r} is not finite.")
+
+    @classmethod
+    def _validate_checkpoint_state_compatibility(
+        cls,
+        model: nn.Module,
+        state_dict: dict,
+    ) -> None:
+        """Validate checkpoint keys, shapes, and destination conversions."""
+        cls._require_finite_state_dict(state_dict)
+        destination_state = model.state_dict()
+        missing = [name for name in destination_state if name not in state_dict]
+        unexpected = [name for name in state_dict if name not in destination_state]
+        if missing or unexpected:
+            details = []
+            if missing:
+                details.append(f"missing keys: {missing}")
+            if unexpected:
+                details.append(f"unexpected keys: {unexpected}")
+            raise RuntimeError(
+                "Checkpoint state_dict is incompatible with the model ("
+                + "; ".join(details)
+                + ")."
+            )
+
+        for name, destination in destination_state.items():
+            stored = state_dict[name]
+            if tuple(stored.shape) != tuple(destination.shape):
+                raise RuntimeError(
+                    f"Checkpoint state entry {name!r} has shape "
+                    f"{tuple(stored.shape)}, but the model expects "
+                    f"{tuple(destination.shape)}."
+                )
+            try:
+                converted = stored.to(
+                    device=destination.device,
+                    dtype=destination.dtype,
+                )
+            except (RuntimeError, TypeError) as error:
+                raise RuntimeError(
+                    f"Checkpoint state entry {name!r} cannot be converted "
+                    f"to {destination.dtype} on {destination.device}."
+                ) from error
+            if (
+                converted.is_floating_point() or converted.is_complex()
+            ) and not torch.isfinite(converted).all():
+                raise FloatingPointError(
+                    f"Checkpoint state entry {name!r} is not finite after "
+                    f"conversion to {destination.dtype}."
+                )
+
+    @staticmethod
+    def _validate_checkpoint_structure(checkpoint: dict) -> None:
+        """Require the mapping structure used by PyTorch checkpoints."""
+        if not isinstance(checkpoint, dict):
+            raise TypeError("The PyTorch checkpoint must be a dictionary.")
+
+    def _stage_checkpoint(self, checkpoint: dict):
+        """Load checkpoint data into an isolated copy of this instance."""
+        self._validate_checkpoint_structure(checkpoint)
+        staged = object.__new__(self.__class__)
+        staged.__dict__ = copy.deepcopy(self.__dict__)
+        checkpoint_config_names = []
+
+        if staged.model is None:
+            build_input_shape = checkpoint.get("build_input_shape")
+            if build_input_shape is None:
+                raise ValueError(
+                    "This legacy checkpoint does not include build_input_shape. "
+                    "Build the model manually before loading it."
+                )
+
+            init_config = checkpoint.get("init_config", {})
+            if not isinstance(init_config, dict):
+                raise TypeError("Checkpoint init_config must be a dictionary.")
+            for name, value in init_config.items():
+                if hasattr(staged, name):
+                    setattr(staged, name, copy.deepcopy(value))
+                    checkpoint_config_names.append(name)
+
+            try:
+                staged._build_input_shape = tuple(build_input_shape)
+            except TypeError as error:
+                raise TypeError(
+                    "Checkpoint build_input_shape must be an iterable shape."
+                ) from error
+            staged.model = staged._build_model(staged._build_input_shape)
+            staged.model = staged.model.to(staged.device)
+
+        checkpoint_state = checkpoint.get("model_state_dict")
+        staged._validate_checkpoint_state_compatibility(
+            staged.model,
+            checkpoint_state,
+        )
+        staged.model.load_state_dict(checkpoint_state)
+        staged._require_finite_parameters()
+        staged.is_fitted = checkpoint.get("is_fitted", False)
+        if staged._build_input_shape is None:
+            shape = checkpoint.get("build_input_shape")
+            if shape is not None:
+                staged._build_input_shape = tuple(shape)
+
+        return staged, checkpoint_config_names
+
+    def _commit_staged_checkpoint(
+        self,
+        staged,
+        checkpoint_config_names: list[str],
+    ) -> None:
+        """Commit an already validated staged checkpoint atomically."""
+        for name in checkpoint_config_names:
+            setattr(self, name, getattr(staged, name))
+        self.model = staged.model
+        self._build_input_shape = staged._build_input_shape
+        self.is_fitted = staged.is_fitted
+
+    @staticmethod
+    def _loss_to_sample_total(
+        loss: torch.Tensor,
+        batch_sample_count: int,
+        criterion: nn.Module,
+    ) -> float:
+        """Convert a scalar batch loss to a sample-total contribution.
+
+        Mean-reduced and custom scalar criteria are treated as batch means.
+        Sum-reduced criteria are already totals. Epoch histories are then
+        divided by the number of samples, so short final batches are weighted
+        correctly.
+        """
+        reduction = getattr(criterion, "reduction", None)
+        value = float(loss.item())
+        if reduction == "sum":
+            return value
+        return value * batch_sample_count
+
     def fit(
         self,
         X: np.ndarray,
@@ -244,40 +538,8 @@ def fit(
         verbose: int = 1,
         **kwargs,
     ) -> dict[str, list]:
-        """
-        Fit the model.
-
-        Parameters
-        ----------
-        X : np.ndarray
-            Training data.
-        y : np.ndarray, optional
-            Target data. If None, assumes autoencoder (X is target). Default is None.
-        validation_split : float, optional
-            Fraction of data to use for validation. Default is 0.2.
-        epochs : int, optional
-            Maximum number of epochs. Default is 500.
-        batch_size : int, optional
-            Batch size. Default is 64.
-        learning_rate : float, optional
-            Learning rate. Default is 1e-3.
-        optimizer : torch.optim.Optimizer, optional
-            Optimizer to use. If None, uses Adam. Default is None.
-        criterion : torch.nn.Module, optional
-            Loss function. If None, uses MSE. Default is None.
-        patience : int, optional
-            Early stopping patience. Default is 20.
-        verbose : int, optional
-            Verbosity level. Default is 1.
-        **kwargs
-            Additional keyword arguments for model building.
-
-        Returns
-        -------
-        dict[str, list]
-            Training history with 'train_loss' and 'val_loss' keys.
-        """
-
+        """Fit a reconstruction model with finite, sample-weighted losses."""
+        learning_rate = self._validate_learning_rate(learning_rate)
         if not isinstance(X, np.ndarray):
             raise TypeError("X must be a NumPy array.")
         if y is None:
@@ -292,131 +554,137 @@ def fit(
             patience,
         )
         self._validate_or_set_build_input_shape(tuple(X.shape))
+        self.is_fitted = False
 
         if self.model is None:
             self.model = self._build_model(X.shape, **kwargs)
             self.model = self.model.to(self.device)
 
         avoid_singleton = self._requires_non_singleton_training_batches()
-
         if optimizer is None:
             optimizer = torch.optim.Adam(self.model.parameters(), lr=learning_rate)
-
         if criterion is None:
             criterion = nn.MSELoss()
 
-        # Train/validation split
-        n_samples = len(X)
-        idx = np.arange(n_samples)
-        np.random.shuffle(idx)
-        split = int((1 - validation_split) * n_samples)
-        train_idx, val_idx = idx[:split], idx[split:]
-        Xtr, Xval = X[train_idx], X[val_idx]
-
-        ytr, yval = y[train_idx], y[val_idx]
+        indices = np.arange(len(X))
+        np.random.shuffle(indices)
+        split = int((1 - validation_split) * len(X))
+        train_indices, validation_indices = indices[:split], indices[split:]
 
-        # Convert to tensors
-        Xtr_tensor = torch.FloatTensor(Xtr).to(self.device)
-        Xval_tensor = torch.FloatTensor(Xval).to(self.device)
-        ytr_tensor = torch.FloatTensor(ytr).to(self.device)
-        yval_tensor = torch.FloatTensor(yval).to(self.device)
+        X_train = torch.as_tensor(
+            X[train_indices], dtype=torch.float32, device=self.device
+        )
+        y_train = torch.as_tensor(
+            y[train_indices], dtype=torch.float32, device=self.device
+        )
+        X_validation = torch.as_tensor(
+            X[validation_indices], dtype=torch.float32, device=self.device
+        )
+        y_validation = torch.as_tensor(
+            y[validation_indices], dtype=torch.float32, device=self.device
+        )
 
         history = {"train_loss": [], "val_loss": []}
-        best_val_loss = float("inf")
+        best_validation_loss = float("inf")
         patience_counter = 0
         best_model_state = None
 
-        # Create progress bar if verbose > 0
-        use_progress_bar = verbose > 0
         epoch_range = range(epochs)
-        pbar = None
-        if use_progress_bar:
-            pbar = tqdm(epoch_range, desc="Training", unit="epoch")
-            epoch_range = pbar
+        progress_bar = None
+        if verbose > 0:
+            progress_bar = tqdm(epoch_range, desc="Training", unit="epoch")
+            epoch_range = progress_bar
 
         for epoch in epoch_range:
-            # Training
             self.model.train()
-            train_loss = 0.0
-            train_slices = self._batch_slices(
-                len(Xtr),
+            train_total = 0.0
+            train_sample_count = 0
+            for start, stop in self._batch_slices(
+                len(X_train),
                 batch_size,
                 avoid_singleton=avoid_singleton,
-            )
-            n_batches = len(train_slices)
-
-            for start, stop in train_slices:
-                batch_X = Xtr_tensor[start:stop]
-                batch_y = ytr_tensor[start:stop]
+            ):
+                batch_X = X_train[start:stop]
+                batch_y = y_train[start:stop]
+                current_batch_size = stop - start
 
                 optimizer.zero_grad()
                 output = self.model(batch_X)
+                self._require_matching_output_shape(output, batch_y, "Training")
+                self._require_finite_tensor(output, "Training output")
+                self._require_finite_buffers()
                 loss = criterion(output, batch_y)
                 self._require_scalar_loss(loss)
+                self._require_finite_loss(loss, "Training")
                 loss.backward()
+                self._require_finite_gradients()
                 optimizer.step()
+                self._require_finite_parameters()
 
-                train_loss += loss.item()
+                train_total += self._loss_to_sample_total(
+                    loss,
+                    current_batch_size,
+                    criterion,
+                )
+                train_sample_count += current_batch_size
 
-            train_loss /= n_batches
+            train_loss = train_total / train_sample_count
             history["train_loss"].append(train_loss)
 
-            # Validation
             self.model.eval()
-            val_loss = 0.0
+            validation_total = 0.0
+            validation_sample_count = 0
             with torch.no_grad():
-                val_slices = self._batch_slices(
-                    len(Xval),
+                for start, stop in self._batch_slices(
+                    len(X_validation),
                     batch_size,
-                    avoid_singleton=False,
-                )
-                n_val_batches = len(val_slices)
-                for start, stop in val_slices:
-                    batch_X = Xval_tensor[start:stop]
-                    batch_y = yval_tensor[start:stop]
-
+                ):
+                    batch_X = X_validation[start:stop]
+                    batch_y = y_validation[start:stop]
+                    current_batch_size = stop - start
                     output = self.model(batch_X)
+                    self._require_matching_output_shape(output, batch_y, "Validation")
+                    self._require_finite_tensor(output, "Validation output")
+                    self._require_finite_parameters()
                     loss = criterion(output, batch_y)
                     self._require_scalar_loss(loss)
-                    val_loss += loss.item()
-
-                val_loss /= n_val_batches
-                history["val_loss"].append(val_loss)
-
-            # Early stopping
-            if val_loss < best_val_loss:
-                best_val_loss = val_loss
+                    self._require_finite_loss(loss, "Validation")
+                    validation_total += self._loss_to_sample_total(
+                        loss,
+                        current_batch_size,
+                        criterion,
+                    )
+                    validation_sample_count += current_batch_size
+
+            validation_loss = validation_total / validation_sample_count
+            history["val_loss"].append(validation_loss)
+
+            if validation_loss < best_validation_loss:
+                best_validation_loss = validation_loss
                 patience_counter = 0
                 best_model_state = copy.deepcopy(self.model.state_dict())
             else:
                 patience_counter += 1
                 if patience_counter >= patience:
-                    if verbose > 0:
-                        if pbar is not None:
-                            pbar.set_postfix_str(f"Early stopping at epoch {epoch + 1}")
-                        self.logger.info(f"Early stopping at epoch {epoch + 1}")
+                    if progress_bar is not None:
+                        progress_bar.set_postfix_str(
+                            f"Early stopping at epoch {epoch + 1}"
+                        )
                     break
 
-            # Update progress bar with current losses
-            if pbar is not None:
-                pbar.set_postfix_str(
+            if progress_bar is not None:
+                progress_bar.set_postfix_str(
                     f"Train Loss: {train_loss:.6f}, "
-                    f"Val Loss: {val_loss:.6f}, "
+                    f"Val Loss: {validation_loss:.6f}, "
                     f"Patience: {patience_counter}/{patience}"
                 )
-            elif verbose > 0 and (epoch + 1) % max(1, epochs // 10) == 0:
-                self.logger.info(
-                    f"Epoch {epoch + 1}/{epochs} - "
-                    f"Train Loss: {train_loss:.6f}, "
-                    f"Val Loss: {val_loss:.6f}"
-                )
-
-        # Restore best model
-        if best_model_state is not None:
-            self.model.load_state_dict(best_model_state)
 
+        if best_model_state is None:
+            raise FloatingPointError(
+                "Training completed without a finite validation loss."
+            )
+        self.model.load_state_dict(best_model_state)
         self.is_fitted = True
-
         return history
 
     def predict(
@@ -447,7 +715,7 @@ def predict(
 
         if not self.is_fitted or self.model is None:
             raise ValueError("Model must be fitted before prediction.")
-        self._validate_inference_inputs(X, batch_size)
+        self._validate_inference_inputs(X, batch_size, check_expected_shape=True)
 
         self.model.eval()
         X_tensor = torch.FloatTensor(X).to(self.device)
@@ -465,6 +733,8 @@ def predict(
             for i in batch_range:
                 batch_X = X_tensor[i : i + batch_size]
                 output = self.model(batch_X)
+                self._require_finite_tensor(output, "Prediction output")
+                self._require_finite_parameters()
                 predictions.append(output.cpu().numpy())
 
         return np.concatenate(predictions, axis=0)
@@ -497,7 +767,7 @@ def encode(
 
         if not self.is_fitted or self.model is None:
             raise ValueError("Model must be fitted before encoding.")
-        self._validate_inference_inputs(X, batch_size)
+        self._validate_inference_inputs(X, batch_size, check_expected_shape=True)
 
         # Check if model has encode_forward method
         if not hasattr(self.model, "encode_forward"):
@@ -522,11 +792,12 @@ def encode(
             for i in batch_range:
                 batch_X = X_tensor[i : i + batch_size]
                 encoding = self.model.encode_forward(batch_X)
+                self._require_finite_tensor(encoding, "Encoding output")
+                self._require_finite_parameters()
                 encodings.append(encoding.cpu().numpy())
 
         return np.concatenate(encodings, axis=0)
 
-
     def decode(
         self,
         Z: np.ndarray,
@@ -544,7 +815,7 @@ def decode(
         Z = np.asarray(Z)
         if Z.ndim == 1:
             Z = Z[None, :]
-        self._validate_inference_inputs(Z, batch_size, name="Z")
+        self._validate_latent_inputs(Z, batch_size)
 
         self.model.eval()
         Z_tensor = torch.as_tensor(Z, dtype=torch.float32, device=self.device)
@@ -562,9 +833,9 @@ def decode(
 
         with torch.no_grad():
             for start in batch_range:
-                output = self.model.decode_forward(
-                    Z_tensor[start : start + batch_size]
-                )
+                output = self.model.decode_forward(Z_tensor[start : start + batch_size])
+                self._require_finite_tensor(output, "Decoding output")
+                self._require_finite_parameters()
                 outputs.append(output.cpu().numpy())
 
         return np.concatenate(outputs, axis=0)
@@ -580,8 +851,11 @@ def reconstruction_error(
         eps: float = 0.0,
     ):
         """Compute reconstruction error with the shared metrics module."""
-        self._validate_inference_inputs(X, batch_size)
+        eps = _validate_eps(eps)
+        self._validate_inference_inputs(X, batch_size, check_expected_shape=True)
         target = self._get_reconstruction_target(X) if y is None else y
+        self._validate_target_shape(X, target)
+        self._validate_finite_array(target, "y")
         prediction = self.predict(X, batch_size=batch_size, verbose=verbose)
         return reconstruction_error_metric(
             target,
@@ -601,8 +875,11 @@ def evaluate_reconstruction(
         eps: float = 0.0,
     ) -> dict[str, float]:
         """Return summary statistics for reconstruction error."""
-        self._validate_inference_inputs(X, batch_size)
+        eps = _validate_eps(eps)
+        self._validate_inference_inputs(X, batch_size, check_expected_shape=True)
         target = self._get_reconstruction_target(X) if y is None else y
+        self._validate_target_shape(X, target)
+        self._validate_finite_array(target, "y")
         prediction = self.predict(X, batch_size=batch_size, verbose=verbose)
         return evaluate_reconstruction_metric(
             target,
@@ -653,6 +930,7 @@ def load_pytorch_model(
             map_location=map_location,
             **kwargs,
         )
+        self._validate_checkpoint_structure(checkpoint)
         checkpoint_class = checkpoint.get("model_class")
         if checkpoint_class and checkpoint_class != self.__class__.__name__:
             raise ValueError(
@@ -660,29 +938,8 @@ def load_pytorch_model(
                 f"not {self.__class__.__name__}."
             )
 
-        if self.model is None:
-            build_input_shape = checkpoint.get("build_input_shape")
-            if build_input_shape is None:
-                raise ValueError(
-                    "This legacy checkpoint does not include build_input_shape. "
-                    "Build the model manually before loading it."
-                )
-
-            init_config = checkpoint.get("init_config", {})
-            for name, value in init_config.items():
-                if hasattr(self, name):
-                    setattr(self, name, value)
-
-            self._build_input_shape = tuple(build_input_shape)
-            self.model = self._build_model(self._build_input_shape)
-            self.model = self.model.to(self.device)
-
-        self.model.load_state_dict(checkpoint["model_state_dict"])
-        self.is_fitted = checkpoint.get("is_fitted", False)
-        if self._build_input_shape is None:
-            shape = checkpoint.get("build_input_shape")
-            if shape is not None:
-                self._build_input_shape = tuple(shape)
+        staged, checkpoint_config_names = self._stage_checkpoint(checkpoint)
+        self._commit_staged_checkpoint(staged, checkpoint_config_names)
         self.logger.info(f"PyTorch model loaded from {model_path}")
         return self
 
@@ -700,6 +957,7 @@ def from_pytorch_model(
             map_location=map_location,
             **kwargs,
         )
+        cls._validate_checkpoint_structure(checkpoint)
         checkpoint_class = checkpoint.get("model_class")
         if checkpoint_class and checkpoint_class != cls.__name__:
             raise ValueError(
@@ -714,10 +972,8 @@ def from_pytorch_model(
                 "init_config and build_input_shape."
             )
 
-        instance = cls(device=device, **dict(init_config))
-        instance._build_input_shape = tuple(build_input_shape)
-        instance.model = instance._build_model(instance._build_input_shape)
-        instance.model = instance.model.to(instance.device)
-        instance.model.load_state_dict(checkpoint["model_state_dict"])
-        instance.is_fitted = checkpoint.get("is_fitted", False)
-        return instance
+        if not isinstance(init_config, dict):
+            raise TypeError("Checkpoint init_config must be a dictionary.")
+        instance = cls(device=device, **copy.deepcopy(init_config))
+        staged, _ = instance._stage_checkpoint(checkpoint)
+        return staged
diff --git a/bluemath_tk/deeplearning/autoencoders.py b/bluemath_tk/deeplearning/autoencoders.py
index 7b09d58..2705a15 100644
--- a/bluemath_tk/deeplearning/autoencoders.py
+++ b/bluemath_tk/deeplearning/autoencoders.py
@@ -6,11 +6,13 @@
 This module contains the following autoencoders:
 - StandardAutoencoder
 - OrthogonalAutoencoder
+- VariationalAutoencoder
 - LSTMAutoencoder
 - CNNAutoencoder
 - VisionTransformerAutoencoder
 - ConvLSTMAutoencoder
 - HybridConvLSTMTransformerAutoencoder
+- SpatialTokenConvLSTMTransformerAutoencoder
 
 Each autoencoder is a subclass of BaseDeepLearningModel and implements the following methods:
 - fit(X, y=None, epochs=10, batch_size=32, verbose=1)
@@ -18,9 +20,20 @@
 - encode(X)
 - decode(X)
 - evaluate(X)
+
+Limitations
+-----------
+Inputs are converted to float32. Complete training and validation splits are
+currently transferred to the selected device instead of being streamed one
+batch at a time, and full-resolution ConvLSTM activations can dominate memory.
+A variational model with beta=0 is not a prior-matched generative model.
+Latent dimensional compression is neither entropy coding nor a deployable,
+bitrate-controlled codec. CUDA RNG preservation is implemented in the test
+suite but still requires validation on CUDA-capable hardware.
 """
 
 import copy
+import math
 from typing import Dict, Optional, Tuple
 
 import numpy as np
@@ -38,6 +51,58 @@
     TimePositionalEncoding,
     Unpatchify,
 )
+from .spatiotemporal_autoencoders import (
+    SpatialTokenConvLSTMTransformerAutoencoder as _SpatialTokenAutoencoder,
+)
+from .variational_autoencoders import (
+    VariationalAutoencoder as _VariationalAutoencoder,
+)
+
+SpatialTokenConvLSTMTransformerAutoencoder = _SpatialTokenAutoencoder
+VariationalAutoencoder = _VariationalAutoencoder
+
+
+__all__ = [
+    "StandardAutoencoder",
+    "OrthogonalAutoencoder",
+    "LSTMAutoencoder",
+    "CNNAutoencoder",
+    "VisionTransformerAutoencoder",
+    "ConvLSTMAutoencoder",
+    "HybridConvLSTMTransformerAutoencoder",
+    "VariationalAutoencoder",
+    "SpatialTokenConvLSTMTransformerAutoencoder",
+]
+
+
+def _validate_positive_integer(name: str, value: int) -> int:
+    if not isinstance(value, int) or isinstance(value, bool) or value < 1:
+        raise ValueError(f"{name} must be a positive integer.")
+    return value
+
+
+def _validate_positive_integer_sequence(
+    name: str,
+    values,
+    expected_length: int | None = None,
+) -> list[int]:
+    if not isinstance(values, (list, tuple)) or not values:
+        raise ValueError(f"{name} must contain positive integers.")
+    if expected_length is not None and len(values) != expected_length:
+        raise ValueError(f"{name} must contain exactly {expected_length} values.")
+    validated = [_validate_positive_integer(f"{name} entry", value) for value in values]
+    return validated
+
+
+def _validate_nonnegative_number(name: str, value: float) -> float:
+    if (
+        not isinstance(value, (int, float))
+        or isinstance(value, bool)
+        or not math.isfinite(float(value))
+        or value < 0
+    ):
+        raise ValueError(f"{name} must be a finite non-negative number.")
+    return float(value)
 
 
 class StandardAutoencoder(BaseDeepLearningModel):
@@ -84,10 +149,12 @@ def __init__(
         device: Optional[torch.device] = None,
         **kwargs,
     ):
+        self.k = _validate_positive_integer("k", k)
         if hidden_dims is None:
             hidden_dims = [512, 256, 128, 64]
-        self.hidden_dims = hidden_dims
-        self.k = k
+        self.hidden_dims = _validate_positive_integer_sequence(
+            "hidden_dims", hidden_dims
+        )
         super().__init__(device=device, **kwargs)
 
     def _build_model(self, input_shape: Tuple, **kwargs) -> nn.Module:
@@ -146,13 +213,14 @@ def encode_forward(self, x):
                     x = x.unsqueeze(0)
                 return self.encoder(x)
 
-
             def decode_forward(self, z):
                 """Decode latent vectors to the original sample shape."""
                 x_recon = self.decoder(z)
                 return x_recon.view(x_recon.size(0), *self.sample_shape)
 
-        return StandardAutoencoderModel(n_features, self.hidden_dims, self.k, sample_shape)
+        return StandardAutoencoderModel(
+            n_features, self.hidden_dims, self.k, sample_shape
+        )
 
 
 class OrthogonalAutoencoder(BaseDeepLearningModel):
@@ -204,12 +272,14 @@ def __init__(
         device: Optional[torch.device] = None,
         **kwargs,
     ):
+        self.k = _validate_positive_integer("k", k)
         if hidden_dims is None:
             hidden_dims = [512, 256, 128, 64]
-        self.hidden_dims = hidden_dims
-        self.k = k
-        self.lambda_W = lambda_W
-        self.lambda_Z = lambda_Z
+        self.hidden_dims = _validate_positive_integer_sequence(
+            "hidden_dims", hidden_dims
+        )
+        self.lambda_W = _validate_nonnegative_number("lambda_W", lambda_W)
+        self.lambda_Z = _validate_nonnegative_number("lambda_Z", lambda_Z)
         super().__init__(device=device, **kwargs)
 
     def _build_model(self, input_shape: Tuple, **kwargs) -> nn.Module:
@@ -302,7 +372,6 @@ def get_regularization_losses(self):
                 decorr_loss = getattr(self.latent_decorr, "_loss", None)
                 return ortho_loss, decorr_loss
 
-
             def decode_forward(self, z):
                 """Decode latent vectors to the original sample shape."""
                 x_recon = self.decoder(z)
@@ -331,12 +400,8 @@ def fit(
         verbose: int = 1,
         **kwargs,
     ) -> Dict[str, list]:
-        """
-        Fit the orthogonal autoencoder with regularization losses.
-
-        This method overrides the base fit() to properly add orthogonality
-        and decorrelation regularization losses during training.
-        """
+        """Fit with orthogonality and latent-decorrelation penalties."""
+        learning_rate = self._validate_learning_rate(learning_rate)
         if not isinstance(X, np.ndarray):
             raise TypeError("X must be a NumPy array.")
         if y is None:
@@ -350,146 +415,153 @@ def fit(
             patience,
         )
         self._validate_or_set_build_input_shape(tuple(X.shape))
+        self.is_fitted = False
 
         if self.model is None:
-            self.model = self._build_model(X.shape, **kwargs)
-            self.model = self.model.to(self.device)
+            self.model = self._build_model(X.shape, **kwargs).to(self.device)
 
         avoid_singleton = self._requires_non_singleton_training_batches()
-
         if optimizer is None:
             optimizer = torch.optim.Adam(self.model.parameters(), lr=learning_rate)
-
         if criterion is None:
             criterion = nn.MSELoss()
 
-        # Train/validation split
-        n_samples = len(X)
-        idx = np.arange(n_samples)
-        np.random.shuffle(idx)
-        split = int((1 - validation_split) * n_samples)
-        train_idx, val_idx = idx[:split], idx[split:]
-        Xtr, Xval = X[train_idx], X[val_idx]
-
-        if y is None:
-            # Autoencoder case
-            ytr, yval = Xtr, Xval
-        else:
-            ytr, yval = y[train_idx], y[val_idx]
-
-        # Convert to tensors
-        Xtr_tensor = torch.FloatTensor(Xtr).to(self.device)
-        Xval_tensor = torch.FloatTensor(Xval).to(self.device)
-        ytr_tensor = torch.FloatTensor(ytr).to(self.device)
-        yval_tensor = torch.FloatTensor(yval).to(self.device)
+        indices = np.arange(len(X))
+        np.random.shuffle(indices)
+        split = int((1 - validation_split) * len(X))
+        train_indices, validation_indices = indices[:split], indices[split:]
+        X_train = torch.as_tensor(
+            X[train_indices], dtype=torch.float32, device=self.device
+        )
+        y_train = torch.as_tensor(
+            y[train_indices], dtype=torch.float32, device=self.device
+        )
+        X_validation = torch.as_tensor(
+            X[validation_indices], dtype=torch.float32, device=self.device
+        )
+        y_validation = torch.as_tensor(
+            y[validation_indices], dtype=torch.float32, device=self.device
+        )
 
         history = {"train_loss": [], "val_loss": []}
-        best_val_loss = float("inf")
+        best_validation_loss = float("inf")
         patience_counter = 0
         best_model_state = None
 
-        # Create progress bar if verbose > 0
-        use_progress_bar = verbose > 0
         epoch_range = range(epochs)
-        pbar = None
-        if use_progress_bar:
-            pbar = tqdm(epoch_range, desc="Training", unit="epoch")
-            epoch_range = pbar
+        progress_bar = None
+        if verbose > 0:
+            progress_bar = tqdm(epoch_range, desc="Training", unit="epoch")
+            epoch_range = progress_bar
 
         for epoch in epoch_range:
-            # Training
             self.model.train()
-            train_loss = 0.0
-            train_slices = self._batch_slices(
-                len(Xtr),
+            train_total = 0.0
+            train_sample_count = 0
+            for start, stop in self._batch_slices(
+                len(X_train),
                 batch_size,
                 avoid_singleton=avoid_singleton,
-            )
-            n_batches = len(train_slices)
-            for start, stop in train_slices:
-                batch_X = Xtr_tensor[start:stop]
-                batch_y = ytr_tensor[start:stop]
+            ):
+                batch_X = X_train[start:stop]
+                batch_y = y_train[start:stop]
+                current_batch_size = stop - start
 
                 optimizer.zero_grad()
                 output = self.model(batch_X)
-                loss = criterion(output, batch_y)
-
-                # Add regularization losses
+                self._require_matching_output_shape(
+                    output, batch_y, "Orthogonal training"
+                )
+                self._require_finite_tensor(output, "Orthogonal training output")
+                self._require_finite_buffers()
+                reconstruction_loss = criterion(output, batch_y)
+                self._require_scalar_loss(reconstruction_loss)
                 ortho_loss, decorr_loss = self.model.get_regularization_losses()
+                regularization_loss = torch.zeros(
+                    (), device=self.device, dtype=reconstruction_loss.dtype
+                )
                 if ortho_loss is not None:
-                    loss = loss + ortho_loss
+                    regularization_loss = regularization_loss + ortho_loss
                 if decorr_loss is not None:
-                    loss = loss + decorr_loss
-
-                self._require_scalar_loss(loss)
+                    regularization_loss = regularization_loss + decorr_loss
+                loss = reconstruction_loss + regularization_loss
+                self._require_finite_loss(loss, "Orthogonal training")
                 loss.backward()
+                self._require_finite_gradients()
                 optimizer.step()
+                self._require_finite_parameters()
 
-                train_loss += loss.item()
+                train_total += self._loss_to_sample_total(
+                    reconstruction_loss,
+                    current_batch_size,
+                    criterion,
+                )
+                train_total += float(regularization_loss.item()) * current_batch_size
+                train_sample_count += current_batch_size
 
-            train_loss /= n_batches
+            train_loss = train_total / train_sample_count
             history["train_loss"].append(train_loss)
 
-            # Validation
             self.model.eval()
-            val_loss = 0.0
+            validation_total = 0.0
+            validation_sample_count = 0
             with torch.no_grad():
-                val_slices = self._batch_slices(
-                    len(Xval),
+                for start, stop in self._batch_slices(
+                    len(X_validation),
                     batch_size,
-                    avoid_singleton=False,
-                )
-                n_val_batches = len(val_slices)
-                for start, stop in val_slices:
-                    batch_X = Xval_tensor[start:stop]
-                    batch_y = yval_tensor[start:stop]
-
+                ):
+                    batch_X = X_validation[start:stop]
+                    batch_y = y_validation[start:stop]
+                    current_batch_size = stop - start
                     output = self.model(batch_X)
-                    loss = criterion(output, batch_y)
-
-                    # Add regularization losses for validation
+                    self._require_matching_output_shape(
+                        output, batch_y, "Orthogonal validation"
+                    )
+                    self._require_finite_tensor(output, "Orthogonal validation output")
+                    self._require_finite_parameters()
+                    reconstruction_loss = criterion(output, batch_y)
+                    self._require_scalar_loss(reconstruction_loss)
                     ortho_loss, decorr_loss = self.model.get_regularization_losses()
+                    regularization_loss = torch.zeros(
+                        (), device=self.device, dtype=reconstruction_loss.dtype
+                    )
                     if ortho_loss is not None:
-                        loss = loss + ortho_loss
+                        regularization_loss = regularization_loss + ortho_loss
                     if decorr_loss is not None:
-                        loss = loss + decorr_loss
-
-                    self._require_scalar_loss(loss)
-                    val_loss += loss.item()
-
-                val_loss /= n_val_batches
-                history["val_loss"].append(val_loss)
+                        regularization_loss = regularization_loss + decorr_loss
+                    loss = reconstruction_loss + regularization_loss
+                    self._require_finite_loss(loss, "Orthogonal validation")
+                    validation_total += self._loss_to_sample_total(
+                        reconstruction_loss,
+                        current_batch_size,
+                        criterion,
+                    )
+                    validation_total += (
+                        float(regularization_loss.item()) * current_batch_size
+                    )
+                    validation_sample_count += current_batch_size
 
-            # Early stopping
-            if val_loss < best_val_loss:
-                best_val_loss = val_loss
+            validation_loss = validation_total / validation_sample_count
+            history["val_loss"].append(validation_loss)
+            if validation_loss < best_validation_loss:
+                best_validation_loss = validation_loss
                 patience_counter = 0
                 best_model_state = copy.deepcopy(self.model.state_dict())
             else:
                 patience_counter += 1
                 if patience_counter >= patience:
-                    if verbose > 0:
-                        if pbar is not None:
-                            pbar.set_postfix_str(f"Early stopping at epoch {epoch + 1}")
-                        self.logger.info(f"Early stopping at epoch {epoch + 1}")
+                    if progress_bar is not None:
+                        progress_bar.set_postfix_str(
+                            f"Early stopping at epoch {epoch + 1}"
+                        )
                     break
 
-            # Update progress bar with current losses
-            if pbar is not None:
-                pbar.set_postfix_str(
-                    f"Train Loss: {train_loss:.6f}, Val Loss: {val_loss:.6f}, Patience: {patience_counter}/{patience}"
-                )
-            elif verbose > 0 and (epoch + 1) % max(1, epochs // 10) == 0:
-                self.logger.info(
-                    f"Epoch {epoch + 1}/{epochs} - Train Loss: {train_loss:.6f}, Val Loss: {val_loss:.6f}"
-                )
-
-        # Restore best model
-        if best_model_state is not None:
-            self.model.load_state_dict(best_model_state)
-
+        if best_model_state is None:
+            raise FloatingPointError(
+                "Training completed without a finite validation loss."
+            )
+        self.model.load_state_dict(best_model_state)
         self.is_fitted = True
-
         return history
 
 
@@ -536,8 +608,10 @@ def __init__(
         device: Optional[torch.device] = None,
         **kwargs,
     ):
-        self.hidden = hidden
-        self.k = k
+        self.k = _validate_positive_integer("k", k)
+        self.hidden = tuple(
+            _validate_positive_integer_sequence("hidden", hidden, expected_length=2)
+        )
         super().__init__(device=device, **kwargs)
 
     def _build_model(self, input_shape: Tuple, **kwargs) -> nn.Module:
@@ -598,13 +672,10 @@ def encode_forward(self, x):
                 z = self.latent(x[:, -1, :])  # Take last timestep
                 return z
 
-
             def decode_forward(self, z):
                 """Decode latent vectors to full temporal sequences."""
                 z_expanded = (
-                    self.latent_to_seq(z)
-                    .unsqueeze(1)
-                    .repeat(1, self.seq_len, 1)
+                    self.latent_to_seq(z).unsqueeze(1).repeat(1, self.seq_len, 1)
                 )
                 x, _ = self.lstm3(z_expanded)
                 x, _ = self.lstm4(x)
@@ -654,7 +725,7 @@ def __init__(
         device: Optional[torch.device] = None,
         **kwargs,
     ):
-        self.k = k
+        self.k = _validate_positive_integer("k", k)
         super().__init__(device=device, **kwargs)
 
     def _build_model(self, input_shape: Tuple, **kwargs) -> nn.Module:
@@ -791,7 +862,6 @@ def encode_forward(self, x):
 
                 return z
 
-
             def decode_forward(self, z):
                 """Decode latent vectors to channels-first spatial grids."""
                 batch_size = z.size(0)
@@ -867,12 +937,16 @@ def __init__(
         device: Optional[torch.device] = None,
         **kwargs,
     ):
-        self.patch_size = patch_size
-        self.d_model = d_model
-        self.depth_enc = depth_enc
-        self.depth_dec = depth_dec
-        self.heads = heads
-        self.k = k
+        self.k = _validate_positive_integer("k", k)
+        self.patch_size = _validate_positive_integer("patch_size", patch_size)
+        self.d_model = _validate_positive_integer("d_model", d_model)
+        if self.d_model < 3:
+            raise ValueError("d_model must be at least 3.")
+        self.depth_enc = _validate_positive_integer("depth_enc", depth_enc)
+        self.depth_dec = _validate_positive_integer("depth_dec", depth_dec)
+        self.heads = _validate_positive_integer("heads", heads)
+        if self.d_model % self.heads != 0:
+            raise ValueError("d_model must be divisible by heads.")
         super().__init__(device=device, **kwargs)
 
     def _build_model(self, input_shape: Tuple, **kwargs) -> nn.Module:
@@ -1061,7 +1135,6 @@ def encode_forward(self, x):
 
                 return z_k
 
-
             def decode_forward(self, z):
                 """Decode latent vectors to channels-first spatial grids."""
                 batch_size = z.size(0)
@@ -1074,9 +1147,7 @@ def decode_forward(self, z):
                 patch_tokens = self.patch_reconstruct(y)
                 reconstruction = self.unpatchify(patch_tokens)
                 if self.pad_h > 0 or self.pad_w > 0:
-                    reconstruction = reconstruction[
-                        :, :, : self.H, : self.W
-                    ]
+                    reconstruction = reconstruction[:, :, : self.H, : self.W]
                 return reconstruction
 
         return ViTAutoencoderModel(
@@ -1095,6 +1166,7 @@ def decode_forward(self, z):
             Pdim,
         )
 
+
 class ConvLSTMAutoencoder(BaseDeepLearningModel):
     """ConvLSTM autoencoder for complete spatiotemporal sequences.
 
@@ -1123,7 +1195,7 @@ def __init__(
                 "reconstruction_mode is no longer supported; "
                 "ConvLSTMAutoencoder always reconstructs the full sequence."
             )
-        self.k = k
+        self.k = _validate_positive_integer("k", k)
         super().__init__(device=device, **kwargs)
 
     def fit(
@@ -1165,8 +1237,7 @@ def _get_reconstruction_target(self, X: np.ndarray) -> np.ndarray:
             raise TypeError("X must be a NumPy array.")
         if X.ndim != 5:
             raise ValueError(
-                "ConvLSTMAutoencoder expects 5D input "
-                "(n_samples, seq_len, C, H, W)."
+                "ConvLSTMAutoencoder expects 5D input (n_samples, seq_len, C, H, W)."
             )
         return X
 
@@ -1188,8 +1259,7 @@ def _build_model(self, input_shape: tuple, **kwargs) -> nn.Module:
         """Build the ConvLSTM encoder and full-sequence decoder."""
         if len(input_shape) != 5:
             raise ValueError(
-                "ConvLSTMAutoencoder expects input shape "
-                "(n_samples, seq_len, C, H, W)."
+                "ConvLSTMAutoencoder expects input shape (n_samples, seq_len, C, H, W)."
             )
 
         seq_len = input_shape[1]
@@ -1275,15 +1345,13 @@ def __init__(self):
             def _validate_input(self, x: torch.Tensor) -> None:
                 if x.dim() != 5:
                     raise ValueError(
-                        "ConvLSTMAutoencoder expects 5D input "
-                        "(B, T, C, H, W)."
+                        "ConvLSTMAutoencoder expects 5D input (B, T, C, H, W)."
                     )
                 sample_shape = tuple(x.shape[1:])
                 expected = (self.seq_len, self.C, self.H, self.W)
                 if sample_shape != expected:
                     raise ValueError(
-                        f"Expected per-sample shape {expected}, "
-                        f"got {sample_shape}."
+                        f"Expected per-sample shape {expected}, got {sample_shape}."
                     )
 
             def _encode(self, x: torch.Tensor) -> torch.Tensor:
@@ -1340,9 +1408,7 @@ def decode_forward(self, z: torch.Tensor) -> torch.Tensor:
                         f"got {tuple(z.shape)}."
                     )
 
-                temporal_input = (
-                    z.unsqueeze(1) + self.decoder_time_embedding
-                )
+                temporal_input = z.unsqueeze(1) + self.decoder_time_embedding
                 temporal_codes, _ = self.temporal_decoder(temporal_input)
                 batch_size = z.size(0)
                 frames = self._decode_spatial(
@@ -1400,10 +1466,16 @@ def __init__(
                 "HybridConvLSTMTransformerAutoencoder always reconstructs "
                 "the full sequence."
             )
-        self.k = k
-        self.d_model = d_model
-        self.n_heads = n_heads
-        self.n_layers = n_layers
+        self.k = _validate_positive_integer("k", k)
+        self.d_model = _validate_positive_integer("d_model", d_model)
+        if self.d_model < 3:
+            raise ValueError("d_model must be at least 3.")
+        self.n_heads = _validate_positive_integer("n_heads", n_heads)
+        self.n_layers = _validate_positive_integer("n_layers", n_layers)
+        if self.d_model % self.n_heads != 0:
+            raise ValueError("d_model must be divisible by n_heads.")
+        if efficient_attention not in {"linear", None}:
+            raise ValueError("efficient_attention must be 'linear' or None.")
         self.efficient_attention = efficient_attention
         super().__init__(device=device, **kwargs)
 
@@ -1626,8 +1698,7 @@ def _validate_input(self, x: torch.Tensor) -> None:
                 expected = (self.seq_len, self.C, self.H, self.W)
                 if sample_shape != expected:
                     raise ValueError(
-                        f"Expected per-sample shape {expected}, "
-                        f"got {sample_shape}."
+                        f"Expected per-sample shape {expected}, got {sample_shape}."
                     )
 
             def _encode(self, x: torch.Tensor) -> torch.Tensor:
@@ -1695,9 +1766,7 @@ def decode_forward(self, z: torch.Tensor) -> torch.Tensor:
                     )
 
                 decoder_tokens = self.latent_to_decoder(z).unsqueeze(1)
-                decoder_tokens = (
-                    decoder_tokens + self.decoder_time_queries
-                )
+                decoder_tokens = decoder_tokens + self.decoder_time_queries
                 decoder_tokens = self.decoder_time_pos_enc(decoder_tokens)
                 decoder_tokens = self._run_blocks(
                     decoder_tokens,
diff --git a/bluemath_tk/deeplearning/layers.py b/bluemath_tk/deeplearning/layers.py
index ace9740..2f99cad 100644
--- a/bluemath_tk/deeplearning/layers.py
+++ b/bluemath_tk/deeplearning/layers.py
@@ -736,9 +736,12 @@ def _extend_for_multilayer(param, num_layers):
 
 class LinearSelfAttention(nn.Module):
     """
-    Softmax-free, Performer-style linear attention on the time axis.
+    Softmax-free linear attention with an ELU+1 feature map.
 
-    Provides O(B * L * D * H) scaling, good for large sequence lengths.
+    Attention contractions scale linearly with sequence length. Including
+    dense projections, total compute is approximately ``O(B * L * D**2)``;
+    accumulated key-value state across all heads is approximately
+    ``O(B * D**2 / H)``.
 
     Parameters
     ----------
@@ -750,7 +753,22 @@ class LinearSelfAttention(nn.Module):
 
     def __init__(self, d_model: int, num_heads: int = 4):
         super().__init__()
-        assert d_model % num_heads == 0
+        for name, value in (
+            ("d_model", d_model),
+            ("num_heads", num_heads),
+        ):
+            if (
+                not isinstance(value, int)
+                or isinstance(value, bool)
+                or value < 1
+            ):
+                raise ValueError(
+                    f"{name} must be a positive integer."
+                )
+        if d_model % num_heads != 0:
+            raise ValueError(
+                "d_model must be divisible by num_heads."
+            )
         self.d_model = d_model
         self.num_heads = num_heads
         self.d_head = d_model // num_heads
@@ -778,6 +796,11 @@ def forward(self, x):
             Output sequences, shape (B, L, D).
         """
         # x: (B, L, D)
+        if x.dim() != 3 or x.size(-1) != self.d_model:
+            raise ValueError(
+                "LinearSelfAttention expects input shape "
+                f"(B, L, {self.d_model}), got {tuple(x.shape)}."
+            )
         B, L, D = x.shape
         Q = self.Wq(x)
         K = self.Wk(x)
diff --git a/bluemath_tk/deeplearning/metrics.py b/bluemath_tk/deeplearning/metrics.py
index b743b4e..95c4c08 100644
--- a/bluemath_tk/deeplearning/metrics.py
+++ b/bluemath_tk/deeplearning/metrics.py
@@ -7,14 +7,15 @@
 
 from __future__ import annotations
 
-from typing import Dict, Literal, Tuple, Union
+import math
+from numbers import Real
+from typing import Literal
 
 import numpy as np
 import torch
 from torch import nn
 
-
-ArrayLike = Union[np.ndarray, torch.Tensor]
+ArrayLike = np.ndarray | torch.Tensor
 MetricName = Literal["mse", "mae", "rmse"]
 ReductionName = Literal["none", "sample", "mean", "sum"]
 
@@ -26,8 +27,12 @@
 ]
 
 
-def _normalise_options(metric: str, reduction: str) -> Tuple[str, str]:
+def _normalise_options(metric: str, reduction: str) -> tuple[str, str]:
     """Validate and normalise metric/reduction names."""
+    if not isinstance(metric, str):
+        raise TypeError("metric must be a string.")
+    if not isinstance(reduction, str):
+        raise TypeError("reduction must be a string.")
     metric = metric.lower()
     reduction = reduction.lower()
 
@@ -44,11 +49,70 @@ def _normalise_options(metric: str, reduction: str) -> Tuple[str, str]:
     return metric, reduction
 
 
+def _validate_eps(eps: float) -> float:
+    """Return a finite, non-negative real scalar epsilon."""
+    if (
+        not isinstance(eps, Real)
+        or isinstance(eps, (bool, np.bool_))
+        or not math.isfinite(float(eps))
+        or eps < 0
+    ):
+        raise ValueError("eps must be a finite, non-negative real scalar.")
+    return float(eps)
+
+
 def _uses_torch(y_true: ArrayLike, y_pred: ArrayLike) -> bool:
     """Return True when at least one input is a PyTorch tensor."""
     return torch.is_tensor(y_true) or torch.is_tensor(y_pred)
 
 
+def _validate_array_like(array: ArrayLike, name: str) -> None:
+    """Require a finite, non-empty real NumPy array or PyTorch tensor."""
+    if isinstance(array, np.ndarray):
+        if array.ndim < 1:
+            raise ValueError(f"{name} must include a sample dimension.")
+        if any(dimension < 1 for dimension in array.shape):
+            raise ValueError(f"{name} dimensions must all be non-empty.")
+        if not np.issubdtype(array.dtype, np.number):
+            raise TypeError(f"{name} must contain numeric values.")
+        if np.issubdtype(array.dtype, np.complexfloating):
+            raise TypeError(f"{name} must contain real-valued data.")
+        if not np.isfinite(array).all():
+            raise ValueError(f"{name} must contain only finite values.")
+        return
+
+    if torch.is_tensor(array):
+        if array.layout != torch.strided or array.is_quantized:
+            raise TypeError(f"{name} must be a dense numeric tensor.")
+        if array.ndim < 1:
+            raise ValueError(f"{name} must include a sample dimension.")
+        if any(dimension < 1 for dimension in array.shape):
+            raise ValueError(f"{name} dimensions must all be non-empty.")
+        if array.dtype == torch.bool:
+            raise TypeError(f"{name} must contain numeric values.")
+        if array.is_complex():
+            raise TypeError(f"{name} must contain real-valued data.")
+        if not torch.isfinite(array).all():
+            raise ValueError(f"{name} must contain only finite values.")
+        return
+
+    raise TypeError(f"{name} must be a NumPy array or PyTorch tensor.")
+
+
+def _as_metric_tensor(array: torch.Tensor) -> torch.Tensor:
+    """Use floating arithmetic for integral tensor metrics."""
+    if array.is_floating_point():
+        return array
+    return array.to(dtype=torch.float64)
+
+
+def _as_metric_array(array: np.ndarray) -> np.ndarray:
+    """Use floating arithmetic for integral NumPy metrics."""
+    if np.issubdtype(array.dtype, np.floating):
+        return array
+    return array.astype(np.float64)
+
+
 def _to_matching_tensor(
     array: ArrayLike,
     reference: torch.Tensor | None = None,
@@ -58,7 +122,17 @@ def _to_matching_tensor(
         return array
 
     if reference is not None:
-        return torch.as_tensor(array, dtype=reference.dtype, device=reference.device)
+        converted = torch.as_tensor(
+            array,
+            dtype=reference.dtype,
+            device=reference.device,
+        )
+        if not torch.isfinite(converted).all():
+            raise FloatingPointError(
+                "Metric input is not finite after conversion to the tensor "
+                "dtype and device."
+            )
+        return converted
 
     return torch.as_tensor(array)
 
@@ -66,7 +140,7 @@ def _to_matching_tensor(
 def _to_numpy(array: ArrayLike) -> np.ndarray:
     """Convert NumPy arrays or tensors to NumPy arrays for summary statistics."""
     if torch.is_tensor(array):
-        return array.detach().cpu().numpy()
+        return array.detach().to(device="cpu", dtype=torch.float64).numpy()
 
     return np.asarray(array)
 
@@ -81,99 +155,235 @@ def _check_same_shape(y_true: ArrayLike, y_pred: ArrayLike) -> None:
         )
 
 
-def _elementwise_error_torch(
-    y_true: torch.Tensor,
-    y_pred: torch.Tensor,
-    metric: str,
-) -> torch.Tensor:
-    """Return elementwise reconstruction error for tensors."""
-    diff = y_pred - y_true
-
-    if metric == "mae":
-        return torch.abs(diff)
-
-    return diff.pow(2)
+def _require_finite_result(result, phase: str) -> None:
+    """Reject non-finite intermediate or final metric arithmetic."""
+    if torch.is_tensor(result):
+        finite = torch.isfinite(result).all()
+    else:
+        finite = np.isfinite(result).all()
+    if not finite:
+        raise FloatingPointError(f"{phase} produced non-finite values.")
 
 
-def _elementwise_error_numpy(
-    y_true: np.ndarray,
-    y_pred: np.ndarray,
-    metric: str,
+def _stable_mean_torch(
+    values: torch.Tensor,
+    dimensions: tuple[int, ...] | None = None,
+) -> torch.Tensor:
+    """Return a max-scaled mean without accumulating unscaled values."""
+    magnitudes = torch.abs(values.detach())
+    if dimensions is None:
+        scale = torch.amax(magnitudes)
+        safe_scale = torch.where(scale > 0, scale, torch.ones_like(scale))
+        return scale * torch.mean(values / safe_scale)
+
+    scale = torch.amax(magnitudes, dim=dimensions, keepdim=True)
+    safe_scale = torch.where(scale > 0, scale, torch.ones_like(scale))
+    normalised_mean = torch.mean(values / safe_scale, dim=dimensions)
+    reduced_scale = scale
+    for dimension in sorted(dimensions, reverse=True):
+        reduced_scale = reduced_scale.squeeze(dimension)
+    return reduced_scale * normalised_mean
+
+
+def _stable_mean_numpy(
+    values: np.ndarray,
+    axes: tuple[int, ...] | None = None,
 ) -> np.ndarray:
-    """Return elementwise reconstruction error for NumPy arrays."""
-    diff = y_pred - y_true
-
-    if metric == "mae":
-        return np.abs(diff)
-
-    return diff**2
+    """Return a max-scaled NumPy mean without unscaled accumulation."""
+    scale = np.max(np.abs(values), axis=axes, keepdims=True)
+    safe_scale = np.where(scale > 0, scale, np.ones_like(scale))
+    normalised_mean = np.mean(values / safe_scale, axis=axes)
+    if axes is None:
+        reduced_scale = np.squeeze(scale)
+    else:
+        reduced_scale = np.squeeze(scale, axis=axes)
+    return reduced_scale * normalised_mean
 
 
 def _reduce_torch(
     elementwise_error: torch.Tensor,
-    metric: str,
     reduction: str,
-    eps: float,
 ) -> torch.Tensor:
-    """Reduce elementwise tensor errors."""
+    """Reduce elementwise MSE or MAE tensor errors."""
     if reduction == "none":
-        if metric == "rmse":
-            return torch.sqrt(elementwise_error + eps)
         return elementwise_error
 
     if elementwise_error.ndim <= 1:
         sample_errors = elementwise_error
     else:
         axes = tuple(range(1, elementwise_error.ndim))
-        sample_errors = elementwise_error.mean(dim=axes)
-
-    if metric == "rmse":
-        sample_errors = torch.sqrt(sample_errors + eps)
+        sample_errors = _stable_mean_torch(elementwise_error, axes)
 
     if reduction == "sample":
         return sample_errors
     if reduction == "mean":
-        return sample_errors.mean()
+        return _stable_mean_torch(sample_errors)
 
     return sample_errors.sum()
 
 
 def _reduce_numpy(
     elementwise_error: np.ndarray,
-    metric: str,
     reduction: str,
-    eps: float,
-) -> Union[np.ndarray, float]:
-    """Reduce elementwise NumPy errors."""
+) -> np.ndarray | float:
+    """Reduce elementwise MSE or MAE NumPy errors."""
     if reduction == "none":
-        if metric == "rmse":
-            return np.sqrt(elementwise_error + eps)
         return elementwise_error
 
     if elementwise_error.ndim <= 1:
         sample_errors = elementwise_error
     else:
         axes = tuple(range(1, elementwise_error.ndim))
-        sample_errors = np.mean(elementwise_error, axis=axes)
+        sample_errors = _stable_mean_numpy(elementwise_error, axes)
+
+    if reduction == "sample":
+        return sample_errors
+    if reduction == "mean":
+        return float(_stable_mean_numpy(sample_errors))
+
+    return float(np.sum(sample_errors))
 
-    if metric == "rmse":
-        sample_errors = np.sqrt(sample_errors + eps)
+
+def _stable_mse_torch(
+    difference: torch.Tensor,
+    dimensions: tuple[int, ...],
+) -> torch.Tensor:
+    """Return a max-scaled tensor MSE over the requested dimensions."""
+    scale = torch.amax(
+        torch.abs(difference.detach()),
+        dim=dimensions,
+        keepdim=True,
+    )
+    safe_scale = torch.where(scale > 0, scale, torch.ones_like(scale))
+    element_count = math.prod(difference.shape[dimension] for dimension in dimensions)
+    scaled_square_sum = torch.sum(
+        ((difference / safe_scale) * difference) / element_count,
+        dim=dimensions,
+    )
+    reduced_scale = scale
+    for dimension in sorted(dimensions, reverse=True):
+        reduced_scale = reduced_scale.squeeze(dimension)
+    return reduced_scale * scaled_square_sum
+
+
+def _mse_torch(difference: torch.Tensor, reduction: str) -> torch.Tensor:
+    """Compute tensor MSE without squaring an unscaled reduction group."""
+    if reduction == "none":
+        return difference.square()
+    if reduction == "mean":
+        return _stable_mse_torch(difference, tuple(range(difference.ndim)))
+
+    if difference.ndim <= 1:
+        sample_errors = difference.square()
+    else:
+        axes = tuple(range(1, difference.ndim))
+        sample_errors = _stable_mse_torch(difference, axes)
 
     if reduction == "sample":
         return sample_errors
+    return sample_errors.sum()
+
+
+def _stable_mse_numpy(
+    difference: np.ndarray,
+    axes: tuple[int, ...],
+) -> np.ndarray:
+    """Return a max-scaled NumPy MSE over the requested axes."""
+    scale = np.max(np.abs(difference), axis=axes, keepdims=True)
+    safe_scale = np.where(scale > 0, scale, np.ones_like(scale))
+    element_count = math.prod(difference.shape[axis] for axis in axes)
+    scaled_square_sum = np.sum(
+        ((difference / safe_scale) * difference) / element_count,
+        axis=axes,
+    )
+    reduced_scale = np.squeeze(scale, axis=axes)
+    return reduced_scale * scaled_square_sum
+
+
+def _mse_numpy(difference: np.ndarray, reduction: str) -> np.ndarray | float:
+    """Compute NumPy MSE without squaring an unscaled reduction group."""
+    if reduction == "none":
+        return difference**2
     if reduction == "mean":
-        return float(np.mean(sample_errors))
+        result = _stable_mse_numpy(difference, tuple(range(difference.ndim)))
+        return float(result)
 
+    if difference.ndim <= 1:
+        sample_errors = difference**2
+    else:
+        axes = tuple(range(1, difference.ndim))
+        sample_errors = _stable_mse_numpy(difference, axes)
+
+    if reduction == "sample":
+        return sample_errors
     return float(np.sum(sample_errors))
 
 
+def _rmse_torch(
+    difference: torch.Tensor,
+    reduction: str,
+    eps: float,
+) -> torch.Tensor:
+    """Compute exact, autograd-safe tensor RMSE reductions."""
+    if reduction == "none" or difference.ndim <= 1:
+        errors = torch.abs(difference)
+    else:
+        flattened = difference.reshape(difference.shape[0], -1)
+        scale = torch.amax(torch.abs(flattened.detach()), dim=1, keepdim=True)
+        nonzero = scale > 0
+        safe_scale = torch.where(nonzero, scale, torch.ones_like(scale))
+        normalised = flattened / safe_scale
+        normalised_mean_square = normalised.square().mean(dim=1)
+
+        nonzero = nonzero.squeeze(1)
+        safe_mean_square = torch.where(
+            nonzero,
+            normalised_mean_square,
+            torch.ones_like(normalised_mean_square),
+        )
+        errors = scale.squeeze(1) * torch.sqrt(safe_mean_square)
+
+    if eps > 0:
+        root_eps = torch.full_like(errors, math.sqrt(eps))
+        errors = torch.hypot(errors, root_eps)
+
+    if reduction in {"none", "sample"}:
+        return errors
+    if reduction == "mean":
+        return _stable_mean_torch(errors)
+    return errors.sum()
+
+
+def _rmse_numpy(
+    difference: np.ndarray,
+    reduction: str,
+    eps: float,
+) -> np.ndarray | float:
+    """Compute numerically stable NumPy RMSE reductions."""
+    if reduction == "none" or difference.ndim <= 1:
+        errors = np.abs(difference)
+    else:
+        flattened = difference.reshape(difference.shape[0], -1)
+        scaled = np.abs(flattened) / math.sqrt(flattened.shape[1])
+        errors = np.hypot.reduce(scaled, axis=1)
+
+    if eps > 0:
+        errors = np.hypot(errors, math.sqrt(eps))
+
+    if reduction in {"none", "sample"}:
+        return errors
+    if reduction == "mean":
+        return float(_stable_mean_numpy(errors))
+    return float(np.sum(errors))
+
+
 def reconstruction_error(
     y_true: ArrayLike,
     y_pred: ArrayLike,
     metric: MetricName = "mse",
     reduction: ReductionName = "mean",
     eps: float = 0.0,
-) -> Union[np.ndarray, torch.Tensor, float]:
+) -> np.ndarray | torch.Tensor | float:
     """Compute reconstruction error between target and prediction.
 
     Parameters
@@ -208,35 +418,47 @@ def reconstruction_error(
     can be used as a differentiable loss.
     """
     metric, reduction = _normalise_options(metric, reduction)
+    eps = _validate_eps(eps)
+    _validate_array_like(y_true, "y_true")
+    _validate_array_like(y_pred, "y_pred")
+    _check_same_shape(y_true, y_pred)
 
     if _uses_torch(y_true, y_pred):
         if torch.is_tensor(y_pred):
-            y_pred_tensor = y_pred
+            y_pred_tensor = _as_metric_tensor(y_pred)
             y_true_tensor = _to_matching_tensor(y_true, reference=y_pred_tensor)
         else:
-            y_true_tensor = y_true
+            y_true_tensor = _as_metric_tensor(y_true)
             y_pred_tensor = _to_matching_tensor(y_pred, reference=y_true_tensor)
 
-        _check_same_shape(y_true_tensor, y_pred_tensor)
-
-        elementwise_error = _elementwise_error_torch(
-            y_true_tensor,
-            y_pred_tensor,
-            metric,
-        )
-        return _reduce_torch(elementwise_error, metric, reduction, eps)
-
-    y_true_array = np.asarray(y_true)
-    y_pred_array = np.asarray(y_pred)
-
-    _check_same_shape(y_true_array, y_pred_array)
-
-    elementwise_error = _elementwise_error_numpy(
-        y_true_array,
-        y_pred_array,
-        metric,
-    )
-    return _reduce_numpy(elementwise_error, metric, reduction, eps)
+        difference = y_pred_tensor - y_true_tensor
+        _require_finite_result(difference, "Metric subtraction")
+        if metric == "rmse":
+            result = _rmse_torch(difference, reduction, eps)
+        elif metric == "mse":
+            result = _mse_torch(difference, reduction)
+        else:
+            elementwise_error = torch.abs(difference)
+            _require_finite_result(elementwise_error, "Metric arithmetic")
+            result = _reduce_torch(elementwise_error, reduction)
+        _require_finite_result(result, "Metric reduction")
+        return result
+
+    y_true_array = _as_metric_array(y_true)
+    y_pred_array = _as_metric_array(y_pred)
+    with np.errstate(over="ignore", invalid="ignore"):
+        difference = y_pred_array - y_true_array
+        _require_finite_result(difference, "Metric subtraction")
+        if metric == "rmse":
+            result = _rmse_numpy(difference, reduction, eps)
+        elif metric == "mse":
+            result = _mse_numpy(difference, reduction)
+        else:
+            elementwise_error = np.abs(difference)
+            _require_finite_result(elementwise_error, "Metric arithmetic")
+            result = _reduce_numpy(elementwise_error, reduction)
+    _require_finite_result(result, "Metric reduction")
+    return result
 
 
 def evaluate_reconstruction(
@@ -244,7 +466,7 @@ def evaluate_reconstruction(
     y_pred: ArrayLike,
     metric: MetricName = "mse",
     eps: float = 0.0,
-) -> Dict[str, float]:
+) -> dict[str, float]:
     """Return summary statistics for per-sample reconstruction error.
 
     Parameters
@@ -265,6 +487,7 @@ def evaluate_reconstruction(
         standard deviation, median, minimum and maximum reconstruction error.
     """
     metric, _ = _normalise_options(metric, "sample")
+    eps = _validate_eps(eps)
 
     sample_errors = reconstruction_error(
         y_true,
@@ -276,15 +499,21 @@ def evaluate_reconstruction(
 
     values = np.ravel(_to_numpy(sample_errors)).astype(float)
 
-    return {
-        "metric": metric,
-        "n_samples": int(values.shape[0]),
-        "mean": float(np.mean(values)),
-        "std": float(np.std(values)),
-        "median": float(np.median(values)),
-        "min": float(np.min(values)),
-        "max": float(np.max(values)),
-    }
+    with np.errstate(over="ignore", invalid="ignore"):
+        summary = {
+            "metric": metric,
+            "n_samples": int(values.shape[0]),
+            "mean": float(_stable_mean_numpy(values)),
+            "std": float(np.std(values)),
+            "median": float(np.median(values)),
+            "min": float(np.min(values)),
+            "max": float(np.max(values)),
+        }
+    if not all(
+        math.isfinite(summary[name]) for name in ("mean", "std", "median", "min", "max")
+    ):
+        raise FloatingPointError("Reconstruction summary produced non-finite values.")
+    return summary
 
 
 class ReconstructionLoss(nn.Module):
@@ -314,7 +543,7 @@ def __init__(
     ):
         super().__init__()
         self.metric, self.reduction = _normalise_options(metric, reduction)
-        self.eps = eps
+        self.eps = _validate_eps(eps)
 
     def forward(self, y_pred: torch.Tensor, y_true: torch.Tensor) -> torch.Tensor:
         """Compute reconstruction loss.
diff --git a/bluemath_tk/deeplearning/spatiotemporal_autoencoders.py b/bluemath_tk/deeplearning/spatiotemporal_autoencoders.py
new file mode 100644
index 0000000..2e3492b
--- /dev/null
+++ b/bluemath_tk/deeplearning/spatiotemporal_autoencoders.py
@@ -0,0 +1,394 @@
+"""Spatially explicit spatiotemporal autoencoders for BlueMath_tk."""
+
+from __future__ import annotations
+
+import numpy as np
+import torch
+import torch.nn as nn
+import torch.nn.functional as functional
+
+from ._base_model import BaseDeepLearningModel
+from .layers import ConvLSTM
+
+
+class _ChannelLayerNorm2d(nn.Module):
+    """Apply LayerNorm over channels for every spatial position."""
+
+    def __init__(self, n_channels: int):
+        super().__init__()
+        self.normalization = nn.LayerNorm(n_channels)
+
+    def forward(self, x: torch.Tensor) -> torch.Tensor:
+        normalized = self.normalization(x.permute(0, 2, 3, 1))
+        return normalized.permute(0, 3, 1, 2)
+
+
+class _FactorizedSpatiotemporalBlock(nn.Module):
+    """Apply temporal attention, spatial attention, and feed-forward updates."""
+
+    def __init__(self, d_model: int, n_heads: int):
+        super().__init__()
+        self.temporal_norm = nn.LayerNorm(d_model)
+        self.temporal_attention = nn.MultiheadAttention(
+            d_model,
+            n_heads,
+            dropout=0.0,
+            batch_first=True,
+        )
+        self.spatial_norm = nn.LayerNorm(d_model)
+        self.spatial_attention = nn.MultiheadAttention(
+            d_model,
+            n_heads,
+            dropout=0.0,
+            batch_first=True,
+        )
+        self.feed_forward_norm = nn.LayerNorm(d_model)
+        self.feed_forward = nn.Sequential(
+            nn.Linear(d_model, d_model * 4),
+            nn.GELU(),
+            nn.Linear(d_model * 4, d_model),
+        )
+
+    def forward(self, x: torch.Tensor) -> torch.Tensor:
+        batch_size, seq_len, n_tokens, d_model = x.shape
+
+        temporal = x.permute(0, 2, 1, 3).reshape(
+            batch_size * n_tokens,
+            seq_len,
+            d_model,
+        )
+        temporal_normalized = self.temporal_norm(temporal)
+        temporal_update, _ = self.temporal_attention(
+            temporal_normalized,
+            temporal_normalized,
+            temporal_normalized,
+            need_weights=False,
+        )
+        temporal = temporal + temporal_update
+        x = temporal.reshape(
+            batch_size,
+            n_tokens,
+            seq_len,
+            d_model,
+        ).permute(0, 2, 1, 3)
+
+        spatial = x.reshape(batch_size * seq_len, n_tokens, d_model)
+        spatial_normalized = self.spatial_norm(spatial)
+        spatial_update, _ = self.spatial_attention(
+            spatial_normalized,
+            spatial_normalized,
+            spatial_normalized,
+            need_weights=False,
+        )
+        spatial = spatial + spatial_update
+        x = spatial.reshape(batch_size, seq_len, n_tokens, d_model)
+
+        return x + self.feed_forward(self.feed_forward_norm(x))
+
+
+class SpatialTokenConvLSTMTransformerAutoencoder(BaseDeepLearningModel):
+    """ConvLSTM-Transformer autoencoder with explicit spatial tokens.
+
+    The model retains multiple spatial tokens at each timestep, applies
+    factorized temporal and spatial attention, compresses the complete window
+    to one vector of size ``k``, and reconstructs the complete input sequence.
+
+    Factorized standard-attention score work scales approximately as
+    ``O(B * d * (S * T**2 + T * S**2))`` for ``S`` spatial tokens and ``T``
+    timesteps, in addition to projection and feed-forward work of
+    ``O(B * T * S * d**2)``.
+    """
+
+    def __init__(
+        self,
+        k: int = 20,
+        spatial_pool_size: tuple[int, int] = (4, 4),
+        d_model: int = 128,
+        n_heads: int = 4,
+        n_layers: int = 2,
+        device: str | torch.device | None = None,
+        **kwargs,
+    ):
+        if not isinstance(k, int) or isinstance(k, bool) or k < 1:
+            raise ValueError("k must be a positive integer.")
+        if (
+            not isinstance(spatial_pool_size, tuple)
+            or len(spatial_pool_size) != 2
+            or any(
+                not isinstance(value, int) or isinstance(value, bool) or value < 1
+                for value in spatial_pool_size
+            )
+        ):
+            raise ValueError(
+                "spatial_pool_size must be a tuple of two positive integers."
+            )
+        if not isinstance(d_model, int) or isinstance(d_model, bool) or d_model < 3:
+            raise ValueError("d_model must be an integer of at least 3.")
+        if not isinstance(n_heads, int) or isinstance(n_heads, bool):
+            raise ValueError("n_heads must be a positive integer.")
+        if n_heads < 1:
+            raise ValueError("n_heads must be a positive integer.")
+        if d_model % n_heads != 0:
+            raise ValueError("d_model must be divisible by n_heads.")
+        if not isinstance(n_layers, int) or isinstance(n_layers, bool):
+            raise ValueError("n_layers must be a positive integer.")
+        if n_layers < 1:
+            raise ValueError("n_layers must be a positive integer.")
+
+        self.k = k
+        self.spatial_pool_size = tuple(spatial_pool_size)
+        self.d_model = d_model
+        self.n_heads = n_heads
+        self.n_layers = n_layers
+        super().__init__(device=device, **kwargs)
+
+    def fit(
+        self,
+        X: np.ndarray,
+        y: np.ndarray | None = None,
+        validation_split: float = 0.2,
+        epochs: int = 500,
+        batch_size: int = 64,
+        learning_rate: float = 1e-3,
+        optimizer: torch.optim.Optimizer | None = None,
+        criterion: nn.Module | None = None,
+        patience: int = 20,
+        verbose: int = 1,
+        **kwargs,
+    ) -> dict[str, list]:
+        """Fit the model to reconstruct the complete input sequence."""
+        if not isinstance(X, np.ndarray):
+            raise TypeError("X must be a NumPy array.")
+        target = self._get_reconstruction_target(X) if y is None else y
+        return super().fit(
+            X,
+            y=target,
+            validation_split=validation_split,
+            epochs=epochs,
+            batch_size=batch_size,
+            learning_rate=learning_rate,
+            optimizer=optimizer,
+            criterion=criterion,
+            patience=patience,
+            verbose=verbose,
+            **kwargs,
+        )
+
+    def _get_reconstruction_target(self, X: np.ndarray) -> np.ndarray:
+        """Use the complete spatiotemporal sequence as the target."""
+        if not isinstance(X, np.ndarray):
+            raise TypeError("X must be a NumPy array.")
+        if X.ndim != 5:
+            raise ValueError(
+                "SpatialTokenConvLSTMTransformerAutoencoder expects 5D input "
+                "(n_samples, seq_len, C, H, W)."
+            )
+        return X
+
+    def _build_model(self, input_shape: tuple, **kwargs) -> nn.Module:
+        """Build the spatial-token encoder and full-sequence decoder."""
+        if len(input_shape) != 5:
+            raise ValueError(
+                "SpatialTokenConvLSTMTransformerAutoencoder expects input "
+                "shape (n_samples, seq_len, C, H, W)."
+            )
+
+        seq_len = input_shape[1]
+        channels, height, width = input_shape[2:]
+        if any(value < 1 for value in (seq_len, channels, height, width)):
+            raise ValueError(
+                "All sequence, channel, and spatial dimensions must be positive."
+            )
+
+        pooled_height, pooled_width = self.spatial_pool_size
+        encoded_height = (height + 3) // 4
+        encoded_width = (width + 3) // 4
+        if pooled_height > encoded_height or pooled_width > encoded_width:
+            raise ValueError(
+                "spatial_pool_size cannot exceed the spatial-encoder output "
+                f"size {(encoded_height, encoded_width)} for input "
+                f"shape {(height, width)}."
+            )
+        n_tokens = pooled_height * pooled_width
+        d_model = self.d_model
+        latent_dim = self.k
+        n_heads = self.n_heads
+        n_layers = self.n_layers
+
+        class SpatialTokenModel(nn.Module):
+            def __init__(self):
+                super().__init__()
+                self.seq_len = seq_len
+                self.channels = channels
+                self.height = height
+                self.width = width
+                self.pooled_height = pooled_height
+                self.pooled_width = pooled_width
+                self.n_tokens = n_tokens
+                self.d_model = d_model
+                self.latent_dim = latent_dim
+
+                self.convlstm1 = ConvLSTM(
+                    input_dim=channels,
+                    hidden_dim=32,
+                    kernel_size=(3, 3),
+                    num_layers=1,
+                    batch_first=True,
+                    return_all_layers=False,
+                )
+                self.convlstm2 = ConvLSTM(
+                    input_dim=32,
+                    hidden_dim=32,
+                    kernel_size=(3, 3),
+                    num_layers=1,
+                    batch_first=True,
+                    return_all_layers=False,
+                )
+                self.spatial_encoder = nn.Sequential(
+                    nn.Conv2d(32, 64, 3, stride=2, padding=1),
+                    _ChannelLayerNorm2d(64),
+                    nn.GELU(),
+                    nn.Conv2d(64, d_model, 3, stride=2, padding=1),
+                    _ChannelLayerNorm2d(d_model),
+                    nn.GELU(),
+                )
+
+                self.encoder_time_embedding = nn.Parameter(
+                    torch.zeros(1, seq_len, 1, d_model)
+                )
+                self.encoder_space_embedding = nn.Parameter(
+                    torch.zeros(1, 1, n_tokens, d_model)
+                )
+                self.encoder_blocks = nn.ModuleList(
+                    [
+                        _FactorizedSpatiotemporalBlock(d_model, n_heads)
+                        for _ in range(n_layers)
+                    ]
+                )
+                self.latent_norm = nn.LayerNorm(d_model)
+                self.latent = nn.Linear(d_model, latent_dim)
+
+                self.latent_to_tokens = nn.Linear(latent_dim, d_model)
+                self.decoder_time_query = nn.Parameter(
+                    torch.zeros(1, seq_len, 1, d_model)
+                )
+                self.decoder_space_query = nn.Parameter(
+                    torch.zeros(1, 1, n_tokens, d_model)
+                )
+                self.decoder_blocks = nn.ModuleList(
+                    [
+                        _FactorizedSpatiotemporalBlock(d_model, n_heads)
+                        for _ in range(n_layers)
+                    ]
+                )
+                self.spatial_decoder = nn.Sequential(
+                    nn.Conv2d(d_model, 64, 3, padding=1),
+                    _ChannelLayerNorm2d(64),
+                    nn.GELU(),
+                    nn.Conv2d(64, 32, 3, padding=1),
+                    _ChannelLayerNorm2d(32),
+                    nn.GELU(),
+                    nn.Conv2d(32, channels, 3, padding=1),
+                )
+
+                for parameter in (
+                    self.encoder_time_embedding,
+                    self.encoder_space_embedding,
+                    self.decoder_time_query,
+                    self.decoder_space_query,
+                ):
+                    nn.init.normal_(parameter, std=0.02)
+
+            def _validate_input(self, x: torch.Tensor) -> None:
+                if x.dim() != 5:
+                    raise ValueError("Expected 5D input with shape (B, T, C, H, W).")
+                expected = (
+                    self.seq_len,
+                    self.channels,
+                    self.height,
+                    self.width,
+                )
+                actual = tuple(x.shape[1:])
+                if actual != expected:
+                    raise ValueError(
+                        f"Expected per-sample shape {expected}, got {actual}."
+                    )
+
+            def _encode_tokens(self, x: torch.Tensor) -> torch.Tensor:
+                self._validate_input(x)
+                batch_size = x.size(0)
+                features, _ = self.convlstm1(x)
+                features, _ = self.convlstm2(features[0])
+                features = features[0]
+
+                features = features.reshape(
+                    batch_size * self.seq_len,
+                    32,
+                    self.height,
+                    self.width,
+                )
+                features = self.spatial_encoder(features)
+                features = functional.adaptive_avg_pool2d(
+                    features,
+                    (self.pooled_height, self.pooled_width),
+                )
+                tokens = features.flatten(2).transpose(1, 2)
+                tokens = tokens.reshape(
+                    batch_size,
+                    self.seq_len,
+                    self.n_tokens,
+                    self.d_model,
+                )
+                tokens = (
+                    tokens + self.encoder_time_embedding + self.encoder_space_embedding
+                )
+                for block in self.encoder_blocks:
+                    tokens = block(tokens)
+                return tokens
+
+            def encode_forward(self, x: torch.Tensor) -> torch.Tensor:
+                tokens = self._encode_tokens(x)
+                pooled = tokens.mean(dim=(1, 2))
+                return self.latent(self.latent_norm(pooled))
+
+            def decode_forward(self, z: torch.Tensor) -> torch.Tensor:
+                if z.dim() != 2 or z.shape[1] != self.latent_dim:
+                    raise ValueError(
+                        f"Latent input must have shape (batch, {self.latent_dim})."
+                    )
+                batch_size = z.size(0)
+                seed = self.latent_to_tokens(z).reshape(
+                    batch_size,
+                    1,
+                    1,
+                    self.d_model,
+                )
+                tokens = seed + self.decoder_time_query + self.decoder_space_query
+                for block in self.decoder_blocks:
+                    tokens = block(tokens)
+
+                features = tokens.permute(0, 1, 3, 2).reshape(
+                    batch_size * self.seq_len,
+                    self.d_model,
+                    self.pooled_height,
+                    self.pooled_width,
+                )
+                features = functional.interpolate(
+                    features,
+                    size=(self.height, self.width),
+                    mode="bilinear",
+                    align_corners=False,
+                )
+                reconstruction = self.spatial_decoder(features)
+                return reconstruction.reshape(
+                    batch_size,
+                    self.seq_len,
+                    self.channels,
+                    self.height,
+                    self.width,
+                )
+
+            def forward(self, x: torch.Tensor) -> torch.Tensor:
+                return self.decode_forward(self.encode_forward(x))
+
+        return SpatialTokenModel()
diff --git a/bluemath_tk/deeplearning/variational_autoencoders.py b/bluemath_tk/deeplearning/variational_autoencoders.py
new file mode 100644
index 0000000..964467f
--- /dev/null
+++ b/bluemath_tk/deeplearning/variational_autoencoders.py
@@ -0,0 +1,631 @@
+"""Variational autoencoders for BlueMath_tk."""
+
+from __future__ import annotations
+
+import copy
+import math
+
+import numpy as np
+import torch
+import torch.nn as nn
+import torch.nn.functional as functional
+from tqdm import tqdm
+
+from ._base_model import BaseDeepLearningModel
+
+
+class VariationalAutoencoder(BaseDeepLearningModel):
+    """Dense variational autoencoder for arbitrary per-sample shapes.
+
+    The encoder parameterizes a diagonal Gaussian posterior. Public
+    :meth:`encode` and deterministic :meth:`predict` use the posterior mean.
+    Stochastic posterior sampling is explicit through ``stochastic=True`` or
+    :meth:`sample_latent`.
+
+    The default objective is an elementwise mean reconstruction loss plus
+    ``beta`` times a KL term summed over latent dimensions and averaged over
+    samples. Therefore, ``beta`` depends on data normalization, per-sample
+    dimensionality, reconstruction-loss scaling, and latent dimension.
+
+    Parameters
+    ----------
+    k : int, optional
+        Number of latent dimensions, by default 20.
+    hidden_dims : list of int, optional
+        Encoder hidden dimensions. The decoder uses the reversed sequence.
+        By default ``[512, 256, 128]``.
+    beta : float, optional
+        Weight applied to the KL-divergence term, by default 1.0.
+    validation_mc_samples : int, optional
+        Posterior samples per validation batch for the stochastic objective
+        used by early stopping, by default 4.
+    device : str or torch.device, optional
+        Device on which to run the model.
+    **kwargs
+        Additional keyword arguments passed to ``BaseDeepLearningModel``.
+    """
+
+    def __init__(
+        self,
+        k: int = 20,
+        hidden_dims: list[int] | None = None,
+        beta: float = 1.0,
+        validation_mc_samples: int = 4,
+        device: str | torch.device | None = None,
+        **kwargs,
+    ):
+        if not isinstance(k, int) or isinstance(k, bool) or k < 1:
+            raise ValueError("k must be a positive integer.")
+        if hidden_dims is None:
+            hidden_dims = [512, 256, 128]
+        if not isinstance(hidden_dims, list) or not hidden_dims:
+            raise ValueError("hidden_dims must be a non-empty list of integers.")
+        if any(
+            not isinstance(dim, int) or isinstance(dim, bool) or dim < 1
+            for dim in hidden_dims
+        ):
+            raise ValueError("Every hidden dimension must be a positive integer.")
+        if (
+            not isinstance(beta, (int, float))
+            or isinstance(beta, bool)
+            or not math.isfinite(float(beta))
+            or beta < 0
+        ):
+            raise ValueError("beta must be a finite non-negative number.")
+        if (
+            not isinstance(validation_mc_samples, int)
+            or isinstance(validation_mc_samples, bool)
+            or validation_mc_samples < 1
+        ):
+            raise ValueError("validation_mc_samples must be a positive integer.")
+
+        self.k = k
+        self.hidden_dims = list(hidden_dims)
+        self.beta = float(beta)
+        self.validation_mc_samples = validation_mc_samples
+        super().__init__(device=device, **kwargs)
+
+    def _build_model(self, input_shape: tuple, **kwargs) -> nn.Module:
+        """Build the encoder, posterior parameterization, and decoder."""
+        if len(input_shape) < 2:
+            raise ValueError(
+                "VariationalAutoencoder requires a leading sample dimension."
+            )
+
+        sample_shape = tuple(input_shape[1:])
+        if any(dimension < 1 for dimension in sample_shape):
+            raise ValueError("Every per-sample dimension must be positive.")
+        n_features = int(np.prod(sample_shape))
+        hidden_dims = tuple(self.hidden_dims)
+        latent_dim = self.k
+
+        class VariationalAutoencoderModel(nn.Module):
+            def __init__(self):
+                super().__init__()
+                self.sample_shape = sample_shape
+                self.n_features = n_features
+                self.latent_dim = latent_dim
+
+                encoder_layers: list[nn.Module] = []
+                previous_dim = n_features
+                for hidden_dim in hidden_dims:
+                    encoder_layers.extend(
+                        [
+                            nn.Linear(previous_dim, hidden_dim),
+                            nn.ReLU(),
+                        ]
+                    )
+                    previous_dim = hidden_dim
+                self.encoder = nn.Sequential(*encoder_layers)
+                self.mu_layer = nn.Linear(previous_dim, latent_dim)
+                self.variance_layer = nn.Linear(previous_dim, latent_dim)
+
+                decoder_layers: list[nn.Module] = []
+                previous_dim = latent_dim
+                for hidden_dim in reversed(hidden_dims):
+                    decoder_layers.extend(
+                        [
+                            nn.Linear(previous_dim, hidden_dim),
+                            nn.ReLU(),
+                        ]
+                    )
+                    previous_dim = hidden_dim
+                decoder_layers.append(nn.Linear(previous_dim, n_features))
+                self.decoder = nn.Sequential(*decoder_layers)
+
+            def _flatten(self, x: torch.Tensor) -> torch.Tensor:
+                if x.dim() < 2:
+                    raise ValueError("Input must include a leading batch dimension.")
+                actual = tuple(x.shape[1:])
+                if actual != self.sample_shape:
+                    raise ValueError(
+                        f"Expected per-sample shape {self.sample_shape}, got {actual}."
+                    )
+                return x.reshape(x.size(0), self.n_features)
+
+            def encode_distribution_forward(
+                self,
+                x: torch.Tensor,
+            ) -> tuple[torch.Tensor, torch.Tensor]:
+                flat = self._flatten(x)
+                hidden = self.encoder(flat)
+                mu = self.mu_layer(hidden)
+                raw_variance = self.variance_layer(hidden)
+                safe_raw = torch.clamp_min(raw_variance, -20.0)
+                central_log_var = torch.log(functional.softplus(safe_raw))
+                log_var = torch.where(
+                    raw_variance < -20.0,
+                    raw_variance,
+                    central_log_var,
+                )
+                return mu, log_var
+
+            @staticmethod
+            def reparameterize(
+                mu: torch.Tensor,
+                log_var: torch.Tensor,
+            ) -> torch.Tensor:
+                standard_deviation = torch.exp(0.5 * log_var)
+                noise = torch.randn_like(standard_deviation)
+                return mu + standard_deviation * noise
+
+            @staticmethod
+            def kl_divergence(
+                mu: torch.Tensor,
+                log_var: torch.Tensor,
+            ) -> torch.Tensor:
+                """Return a numerically stable diagonal-Gaussian KL mean."""
+                mu_stable = mu.to(dtype=torch.float64)
+                log_var_stable = log_var.to(dtype=torch.float64)
+                per_sample = 0.5 * torch.sum(
+                    mu_stable.pow(2) + log_var_stable.exp() - 1.0 - log_var_stable,
+                    dim=1,
+                )
+                return per_sample.mean()
+
+            def decode_forward(self, z: torch.Tensor) -> torch.Tensor:
+                if z.dim() != 2 or z.shape[1] != self.latent_dim:
+                    raise ValueError(
+                        f"Latent input must have shape (batch, {self.latent_dim})."
+                    )
+                reconstruction = self.decoder(z)
+                return reconstruction.reshape(
+                    reconstruction.size(0),
+                    *self.sample_shape,
+                )
+
+            def encode_forward(self, x: torch.Tensor) -> torch.Tensor:
+                mu, _ = self.encode_distribution_forward(x)
+                return mu
+
+            def forward(
+                self,
+                x: torch.Tensor,
+                stochastic: bool | None = None,
+            ) -> torch.Tensor:
+                mu, log_var = self.encode_distribution_forward(x)
+                if stochastic is None:
+                    stochastic = self.training
+                z = self.reparameterize(mu, log_var) if stochastic else mu
+                return self.decode_forward(z)
+
+        return VariationalAutoencoderModel()
+
+    def fit(
+        self,
+        X: np.ndarray,
+        y: np.ndarray | None = None,
+        validation_split: float = 0.2,
+        epochs: int = 500,
+        batch_size: int = 64,
+        learning_rate: float = 1e-3,
+        optimizer: torch.optim.Optimizer | None = None,
+        criterion: nn.Module | None = None,
+        patience: int = 20,
+        verbose: int = 1,
+        **kwargs,
+    ) -> dict[str, list]:
+        """Fit the VAE with stochastic train and validation objectives.
+
+        ``val_loss`` is a Monte Carlo estimate of the same beta-VAE objective
+        used for training and controls early stopping. The separate
+        ``val_deterministic_reconstruction_loss`` reports posterior-mean
+        reconstruction for stable scientific comparison.
+        """
+        learning_rate = self._validate_learning_rate(learning_rate)
+        if not isinstance(X, np.ndarray):
+            raise TypeError("X must be a NumPy array.")
+        if y is None:
+            y = self._get_reconstruction_target(X)
+
+        self._validate_fit_inputs(
+            X,
+            y,
+            validation_split,
+            batch_size,
+            epochs,
+            patience,
+        )
+        self._validate_or_set_build_input_shape(tuple(X.shape))
+        self.is_fitted = False
+
+        if self.model is None:
+            self.model = self._build_model(X.shape, **kwargs).to(self.device)
+
+        if optimizer is None:
+            optimizer = torch.optim.Adam(
+                self.model.parameters(),
+                lr=learning_rate,
+            )
+        if criterion is None:
+            criterion = nn.MSELoss()
+        reduction = getattr(criterion, "reduction", "mean")
+        if reduction not in {"mean", None}:
+            raise ValueError(
+                "VariationalAutoencoder requires a mean-reduced scalar "
+                "reconstruction criterion."
+            )
+
+        indices = np.arange(len(X))
+        np.random.shuffle(indices)
+        split = int((1 - validation_split) * len(X))
+        train_indices = indices[:split]
+        validation_indices = indices[split:]
+
+        X_train = torch.as_tensor(
+            X[train_indices],
+            dtype=torch.float32,
+            device=self.device,
+        )
+        y_train = torch.as_tensor(
+            y[train_indices],
+            dtype=torch.float32,
+            device=self.device,
+        )
+        X_validation = torch.as_tensor(
+            X[validation_indices],
+            dtype=torch.float32,
+            device=self.device,
+        )
+        y_validation = torch.as_tensor(
+            y[validation_indices],
+            dtype=torch.float32,
+            device=self.device,
+        )
+
+        history = {
+            "train_loss": [],
+            "train_reconstruction_loss": [],
+            "train_kl_loss": [],
+            "val_loss": [],
+            "val_reconstruction_loss": [],
+            "val_kl_loss": [],
+            "val_deterministic_reconstruction_loss": [],
+        }
+        best_validation_loss = float("inf")
+        best_model_state = None
+        patience_counter = 0
+
+        epoch_range = range(epochs)
+        progress_bar = None
+        if verbose > 0:
+            progress_bar = tqdm(epoch_range, desc="Training", unit="epoch")
+            epoch_range = progress_bar
+
+        for epoch in epoch_range:
+            self.model.train()
+            train_totals = self._run_vae_epoch(
+                X_train,
+                y_train,
+                batch_size,
+                criterion,
+                optimizer=optimizer,
+                stochastic_samples=1,
+                report_deterministic=False,
+            )
+            history["train_loss"].append(train_totals["loss"])
+            history["train_reconstruction_loss"].append(
+                train_totals["reconstruction_loss"]
+            )
+            history["train_kl_loss"].append(train_totals["kl_loss"])
+
+            self.model.eval()
+            validation_devices: list[int] = []
+            if self.device.type == "cuda":
+                device_index = self.device.index
+                if device_index is None:
+                    device_index = torch.cuda.current_device()
+                validation_devices = [device_index]
+            with torch.random.fork_rng(devices=validation_devices):
+                with torch.no_grad():
+                    validation_totals = self._run_vae_epoch(
+                        X_validation,
+                        y_validation,
+                        batch_size,
+                        criterion,
+                        optimizer=None,
+                        stochastic_samples=self.validation_mc_samples,
+                        report_deterministic=True,
+                    )
+            history["val_loss"].append(validation_totals["loss"])
+            history["val_reconstruction_loss"].append(
+                validation_totals["reconstruction_loss"]
+            )
+            history["val_kl_loss"].append(validation_totals["kl_loss"])
+            history["val_deterministic_reconstruction_loss"].append(
+                validation_totals["deterministic_reconstruction_loss"]
+            )
+
+            validation_loss = validation_totals["loss"]
+            if validation_loss < best_validation_loss:
+                best_validation_loss = validation_loss
+                best_model_state = copy.deepcopy(self.model.state_dict())
+                patience_counter = 0
+            else:
+                patience_counter += 1
+                if patience_counter >= patience:
+                    if progress_bar is not None:
+                        progress_bar.set_postfix_str(
+                            f"Early stopping at epoch {epoch + 1}"
+                        )
+                    break
+
+            if progress_bar is not None:
+                progress_bar.set_postfix_str(
+                    f"Train: {train_totals['loss']:.6f}, "
+                    f"Val: {validation_loss:.6f}, "
+                    f"Patience: {patience_counter}/{patience}"
+                )
+
+        if best_model_state is None:
+            raise FloatingPointError(
+                "Training completed without a finite validation objective."
+            )
+        self.model.load_state_dict(best_model_state)
+        self.is_fitted = True
+        return history
+
+    def _run_vae_epoch(
+        self,
+        X: torch.Tensor,
+        y: torch.Tensor,
+        batch_size: int,
+        criterion: nn.Module,
+        optimizer: torch.optim.Optimizer | None,
+        stochastic_samples: int,
+        report_deterministic: bool,
+    ) -> dict[str, float]:
+        if self.model is None:
+            raise ValueError("Model must be built before training.")
+
+        totals = {
+            "loss": 0.0,
+            "reconstruction_loss": 0.0,
+            "kl_loss": 0.0,
+            "deterministic_reconstruction_loss": 0.0,
+        }
+        total_samples = 0
+
+        for start, stop in self._batch_slices(len(X), batch_size):
+            batch_X = X[start:stop]
+            batch_y = y[start:stop]
+            current_batch_size = stop - start
+            if optimizer is not None:
+                optimizer.zero_grad()
+
+            mu, log_var = self.model.encode_distribution_forward(batch_X)
+            self._require_finite_tensor(mu, "VAE posterior mean")
+            self._require_finite_tensor(log_var, "VAE posterior log variance")
+            reconstruction_losses = []
+            for _ in range(stochastic_samples):
+                z = self.model.reparameterize(mu, log_var)
+                reconstruction = self.model.decode_forward(z)
+                self._require_matching_output_shape(
+                    reconstruction, batch_y, "VAE reconstruction"
+                )
+                self._require_finite_tensor(reconstruction, "VAE reconstruction output")
+                reconstruction_loss = criterion(reconstruction, batch_y)
+                self._require_scalar_loss(reconstruction_loss)
+                self._require_finite_loss(
+                    reconstruction_loss,
+                    "VAE reconstruction",
+                )
+                reconstruction_losses.append(reconstruction_loss)
+
+            mean_reconstruction_loss = torch.stack(reconstruction_losses).mean()
+            kl_loss = self.model.kl_divergence(mu, log_var)
+            self._require_finite_loss(kl_loss, "VAE KL")
+            loss = mean_reconstruction_loss + self.beta * kl_loss
+            self._require_finite_loss(loss, "VAE total")
+
+            deterministic_loss = None
+            if report_deterministic:
+                deterministic = self.model.decode_forward(mu)
+                self._require_matching_output_shape(
+                    deterministic,
+                    batch_y,
+                    "VAE deterministic reconstruction",
+                )
+                self._require_finite_tensor(
+                    deterministic,
+                    "VAE deterministic reconstruction output",
+                )
+                deterministic_loss = criterion(deterministic, batch_y)
+                self._require_scalar_loss(deterministic_loss)
+                self._require_finite_loss(
+                    deterministic_loss,
+                    "VAE deterministic reconstruction",
+                )
+
+            if optimizer is None:
+                self._require_finite_parameters()
+            else:
+                self._require_finite_buffers()
+
+            if optimizer is not None:
+                loss.backward()
+                self._require_finite_gradients()
+                optimizer.step()
+                self._require_finite_parameters()
+
+            totals["loss"] += float(loss.item()) * current_batch_size
+            totals["reconstruction_loss"] += (
+                float(mean_reconstruction_loss.item()) * current_batch_size
+            )
+            totals["kl_loss"] += float(kl_loss.item()) * current_batch_size
+            if deterministic_loss is not None:
+                totals["deterministic_reconstruction_loss"] += (
+                    float(deterministic_loss.item()) * current_batch_size
+                )
+            total_samples += current_batch_size
+
+        return {name: value / total_samples for name, value in totals.items()}
+
+    def predict(
+        self,
+        X: np.ndarray,
+        batch_size: int = 64,
+        verbose: int = 1,
+        stochastic: bool = False,
+    ) -> np.ndarray:
+        """Reconstruct inputs from posterior means or posterior samples."""
+        if not stochastic:
+            return super().predict(
+                X,
+                batch_size=batch_size,
+                verbose=verbose,
+            )
+        if not self.is_fitted or self.model is None:
+            raise ValueError("Model must be fitted before prediction.")
+        self._validate_inference_inputs(
+            X,
+            batch_size,
+            check_expected_shape=True,
+        )
+
+        self.model.eval()
+        X_tensor = torch.as_tensor(
+            X,
+            dtype=torch.float32,
+            device=self.device,
+        )
+        outputs = []
+        with torch.no_grad():
+            for start in range(0, len(X), batch_size):
+                output = self.model(
+                    X_tensor[start : start + batch_size],
+                    stochastic=True,
+                )
+                self._require_finite_tensor(output, "Stochastic prediction output")
+                self._require_finite_parameters()
+                outputs.append(output.cpu().numpy())
+        return np.concatenate(outputs, axis=0)
+
+    def encode_distribution(
+        self,
+        X: np.ndarray,
+        batch_size: int = 64,
+    ) -> tuple[np.ndarray, np.ndarray]:
+        """Return posterior means and log variances."""
+        if not self.is_fitted or self.model is None:
+            raise ValueError("Model must be fitted before encoding.")
+        self._validate_inference_inputs(
+            X,
+            batch_size,
+            check_expected_shape=True,
+        )
+
+        self.model.eval()
+        X_tensor = torch.as_tensor(
+            X,
+            dtype=torch.float32,
+            device=self.device,
+        )
+        means = []
+        log_variances = []
+        with torch.no_grad():
+            for start in range(0, len(X), batch_size):
+                mu, log_var = self.model.encode_distribution_forward(
+                    X_tensor[start : start + batch_size]
+                )
+                self._require_finite_tensor(mu, "Posterior mean output")
+                self._require_finite_tensor(log_var, "Posterior log-variance output")
+                self._require_finite_parameters()
+                means.append(mu.cpu().numpy())
+                log_variances.append(log_var.cpu().numpy())
+        return (
+            np.concatenate(means, axis=0),
+            np.concatenate(log_variances, axis=0),
+        )
+
+    def sample_latent(
+        self,
+        X: np.ndarray,
+        batch_size: int = 64,
+    ) -> np.ndarray:
+        """Draw one posterior latent sample for every input sample."""
+        if not self.is_fitted or self.model is None:
+            raise ValueError("Model must be fitted before sampling latents.")
+        self._validate_inference_inputs(
+            X,
+            batch_size,
+            check_expected_shape=True,
+        )
+
+        self.model.eval()
+        X_tensor = torch.as_tensor(
+            X,
+            dtype=torch.float32,
+            device=self.device,
+        )
+        samples = []
+        with torch.no_grad():
+            for start in range(0, len(X), batch_size):
+                mu, log_var = self.model.encode_distribution_forward(
+                    X_tensor[start : start + batch_size]
+                )
+                self._require_finite_tensor(mu, "Posterior mean output")
+                self._require_finite_tensor(log_var, "Posterior log-variance output")
+                sample = self.model.reparameterize(mu, log_var)
+                self._require_finite_tensor(sample, "Posterior latent sample")
+                self._require_finite_parameters()
+                samples.append(sample.cpu().numpy())
+        return np.concatenate(samples, axis=0)
+
+    def sample(
+        self,
+        n_samples: int,
+        batch_size: int = 64,
+    ) -> np.ndarray:
+        """Decode standard-normal prior samples.
+
+        Prior samples are generatively meaningful only when KL regularization
+        has aligned the learned posterior with the standard-normal prior.
+        """
+        if not isinstance(n_samples, int) or isinstance(n_samples, bool):
+            raise TypeError("n_samples must be an integer.")
+        if n_samples < 1:
+            raise ValueError("n_samples must be at least 1.")
+        if not isinstance(batch_size, int) or isinstance(batch_size, bool):
+            raise TypeError("batch_size must be an integer.")
+        if batch_size < 1:
+            raise ValueError("batch_size must be at least 1.")
+        if not self.is_fitted or self.model is None:
+            raise ValueError("Model must be fitted before sampling.")
+
+        outputs = []
+        self.model.eval()
+        with torch.no_grad():
+            for start in range(0, n_samples, batch_size):
+                current_batch = min(batch_size, n_samples - start)
+                z = torch.randn(
+                    current_batch,
+                    self.k,
+                    dtype=torch.float32,
+                    device=self.device,
+                )
+                output = self.model.decode_forward(z)
+                self._require_finite_tensor(output, "Prior sample output")
+                self._require_finite_parameters()
+                outputs.append(output.cpu().numpy())
+        return np.concatenate(outputs, axis=0)
diff --git a/tests/deeplearning/test_advanced_autoencoder_integration.py b/tests/deeplearning/test_advanced_autoencoder_integration.py
new file mode 100644
index 0000000..3cb0f0e
--- /dev/null
+++ b/tests/deeplearning/test_advanced_autoencoder_integration.py
@@ -0,0 +1,139 @@
+"""Integration tests for the advanced autoencoder public API."""
+
+import numpy as np
+import pytest
+
+torch = pytest.importorskip("torch")
+
+from bluemath_tk.deeplearning import autoencoders  # noqa: E402
+from bluemath_tk.deeplearning.autoencoders import (  # noqa: E402
+    CNNAutoencoder,
+    ConvLSTMAutoencoder,
+    HybridConvLSTMTransformerAutoencoder,
+    LSTMAutoencoder,
+    OrthogonalAutoencoder,
+    SpatialTokenConvLSTMTransformerAutoencoder,
+    StandardAutoencoder,
+    VariationalAutoencoder,
+    VisionTransformerAutoencoder,
+)
+
+
+@pytest.fixture(autouse=True)
+def _set_seed():
+    previous_threads = torch.get_num_threads()
+    numpy_state = np.random.get_state()
+    torch_state = torch.random.get_rng_state()
+    cuda_states = torch.cuda.get_rng_state_all() if torch.cuda.is_available() else None
+    try:
+        np.random.seed(607)
+        torch.manual_seed(607)
+        torch.set_num_threads(1)
+        yield
+    finally:
+        torch.set_num_threads(previous_threads)
+        np.random.set_state(numpy_state)
+        torch.random.set_rng_state(torch_state)
+        if cuda_states is not None:
+            torch.cuda.set_rng_state_all(cuda_states)
+
+
+def test_advanced_autoencoders_are_publicly_exported():
+    assert autoencoders.VariationalAutoencoder is VariationalAutoencoder
+    assert (
+        autoencoders.SpatialTokenConvLSTMTransformerAutoencoder
+        is SpatialTokenConvLSTMTransformerAutoencoder
+    )
+
+
+def test_autoencoder_all_contains_exact_public_classes():
+    expected = {
+        "StandardAutoencoder": StandardAutoencoder,
+        "OrthogonalAutoencoder": OrthogonalAutoencoder,
+        "LSTMAutoencoder": LSTMAutoencoder,
+        "CNNAutoencoder": CNNAutoencoder,
+        "VisionTransformerAutoencoder": VisionTransformerAutoencoder,
+        "ConvLSTMAutoencoder": ConvLSTMAutoencoder,
+        "HybridConvLSTMTransformerAutoencoder": (HybridConvLSTMTransformerAutoencoder),
+        "VariationalAutoencoder": VariationalAutoencoder,
+        "SpatialTokenConvLSTMTransformerAutoencoder": (
+            SpatialTokenConvLSTMTransformerAutoencoder
+        ),
+    }
+
+    assert autoencoders.__all__ == list(expected)
+    for name, intended_class in expected.items():
+        assert getattr(autoencoders, name) is intended_class
+
+
+def test_advanced_autoencoders_follow_common_encode_decode_contract():
+    dense_X = np.random.randn(12, 6).astype("float32")
+    sequence_X = np.random.randn(12, 3, 1, 6, 6).astype("float32")
+
+    models_and_data = [
+        (
+            VariationalAutoencoder(
+                k=3,
+                hidden_dims=[8],
+                beta=0.1,
+                device="cpu",
+            ),
+            dense_X,
+        ),
+        (
+            SpatialTokenConvLSTMTransformerAutoencoder(
+                k=3,
+                spatial_pool_size=(2, 2),
+                d_model=8,
+                n_heads=2,
+                n_layers=1,
+                device="cpu",
+            ),
+            sequence_X,
+        ),
+    ]
+
+    for model, X in models_and_data:
+        model.fit(
+            X,
+            validation_split=0.25,
+            epochs=1,
+            batch_size=4,
+            patience=2,
+            verbose=0,
+        )
+        latent = model.encode(X, verbose=0)
+        prediction = model.predict(X, verbose=0)
+        decoded = model.decode(latent, verbose=0)
+        metrics = model.evaluate(X, metric="rmse", verbose=0)
+
+        assert prediction.shape == X.shape
+        assert decoded.shape == X.shape
+        assert np.allclose(prediction, decoded, rtol=1e-5, atol=1e-6)
+        assert metrics["metric"] == "rmse"
+        assert metrics["n_samples"] == len(X)
+
+
+def test_advanced_checkpoint_class_mismatch_is_rejected(tmp_path):
+    X = np.random.randn(12, 6).astype("float32")
+    vae = VariationalAutoencoder(
+        k=3,
+        hidden_dims=[8],
+        device="cpu",
+    )
+    vae.fit(
+        X,
+        validation_split=0.25,
+        epochs=1,
+        batch_size=4,
+        patience=2,
+        verbose=0,
+    )
+    checkpoint = tmp_path / "vae.pt"
+    vae.save_pytorch_model(checkpoint)
+
+    with pytest.raises(ValueError, match="Checkpoint contains"):
+        SpatialTokenConvLSTMTransformerAutoencoder.from_pytorch_model(
+            checkpoint,
+            device="cpu",
+        )
diff --git a/tests/deeplearning/test_autoencoder_hardening.py b/tests/deeplearning/test_autoencoder_hardening.py
new file mode 100644
index 0000000..c346dec
--- /dev/null
+++ b/tests/deeplearning/test_autoencoder_hardening.py
@@ -0,0 +1,1351 @@
+"""Regression tests for shared autoencoder hardening."""
+
+import copy
+
+import numpy as np
+import pytest
+import torch
+import torch.nn as nn
+
+from bluemath_tk.deeplearning._base_model import BaseDeepLearningModel
+from bluemath_tk.deeplearning.autoencoders import (
+    CNNAutoencoder,
+    ConvLSTMAutoencoder,
+    HybridConvLSTMTransformerAutoencoder,
+    LSTMAutoencoder,
+    OrthogonalAutoencoder,
+    SpatialTokenConvLSTMTransformerAutoencoder,
+    StandardAutoencoder,
+    VariationalAutoencoder,
+    VisionTransformerAutoencoder,
+)
+from bluemath_tk.deeplearning.layers import LinearSelfAttention
+
+
+@pytest.fixture(autouse=True)
+def _set_seed():
+    previous_threads = torch.get_num_threads()
+    numpy_state = np.random.get_state()
+    torch_state = torch.random.get_rng_state()
+    cuda_states = torch.cuda.get_rng_state_all() if torch.cuda.is_available() else None
+    try:
+        np.random.seed(607)
+        torch.manual_seed(607)
+        torch.set_num_threads(1)
+        yield
+    finally:
+        torch.set_num_threads(previous_threads)
+        np.random.set_state(numpy_state)
+        torch.random.set_rng_state(torch_state)
+        if cuda_states is not None:
+            torch.cuda.set_rng_state_all(cuda_states)
+
+
+class _TinyAutoencoder(BaseDeepLearningModel):
+    def __init__(self, device="cpu"):
+        self.mutable_config = {
+            "encoder": {"widths": [2, 1]},
+            "decoder": {"widths": [2, 1]},
+            "flags": {"transactional", "tiny"},
+        }
+        super().__init__(device=device)
+
+    def _build_model(self, input_shape, **kwargs):
+        sample_shape = tuple(input_shape[1:])
+
+        class _Model(nn.Module):
+            def __init__(self):
+                super().__init__()
+                self.anchor = nn.Parameter(torch.zeros(()))
+                self.tail = nn.Parameter(torch.zeros(2))
+                self.register_buffer("running_value", torch.zeros(()))
+
+            def forward(self, x):
+                return torch.zeros_like(x) + self.anchor * 0
+
+            def encode_forward(self, x):
+                return x.reshape(len(x), -1)
+
+            def decode_forward(self, z):
+                return z.reshape(len(z), *sample_shape)
+
+        return _Model()
+
+
+def _find_mutable_objects(value, path=(), active_ids=frozenset()):
+    mutable_objects = {}
+    if isinstance(value, (dict, list, set)):
+        mutable_objects[path] = value
+        value_id = id(value)
+        if value_id in active_ids:
+            return mutable_objects
+        active_ids = active_ids | {value_id}
+
+    if isinstance(value, dict):
+        for key, nested_value in value.items():
+            mutable_objects.update(
+                _find_mutable_objects(
+                    nested_value,
+                    path + (("key", key),),
+                    active_ids,
+                )
+            )
+    elif isinstance(value, (list, tuple)):
+        for index, nested_value in enumerate(value):
+            mutable_objects.update(
+                _find_mutable_objects(
+                    nested_value,
+                    path + (("index", index),),
+                    active_ids,
+                )
+            )
+
+    return mutable_objects
+
+
+def _snapshot_model(model):
+    metadata = {
+        name: value
+        for name, value in model.__dict__.items()
+        if name not in {"_logger", "model"}
+    }
+    mutable_metadata = {}
+    for name, value in metadata.items():
+        mutable_metadata.update(
+            _find_mutable_objects(value, path=(("attribute", name),))
+        )
+
+    return {
+        "model": model.model,
+        "logger": model.__dict__["_logger"],
+        "training": model.model.training if model.model is not None else None,
+        "state": {
+            name: value.detach().clone()
+            for name, value in (
+                model.model.state_dict().items() if model.model is not None else []
+            )
+        },
+        "metadata_keys": set(model.__dict__),
+        "metadata": copy.deepcopy(metadata),
+        "mutable_metadata": mutable_metadata,
+    }
+
+
+def _assert_model_unchanged(model, snapshot):
+    assert set(model.__dict__) == snapshot["metadata_keys"]
+    assert model.model is snapshot["model"]
+    assert model.__dict__["_logger"] is snapshot["logger"]
+    for name, value in snapshot["metadata"].items():
+        assert model.__dict__[name] == value
+
+    mutable_metadata = {}
+    for name in snapshot["metadata"]:
+        mutable_metadata.update(
+            _find_mutable_objects(
+                model.__dict__[name],
+                path=(("attribute", name),),
+            )
+        )
+    assert mutable_metadata.keys() == snapshot["mutable_metadata"].keys()
+    for path, value in snapshot["mutable_metadata"].items():
+        assert mutable_metadata[path] is value
+
+    if model.model is None:
+        assert snapshot["training"] is None
+        assert not snapshot["state"]
+        return
+    assert model.model.training == snapshot["training"]
+    current_state = model.model.state_dict()
+    assert current_state.keys() == snapshot["state"].keys()
+    for name, value in snapshot["state"].items():
+        assert torch.equal(current_state[name], value)
+
+
+def test_unchanged_state_helper_rejects_new_metadata_and_nested_aliases():
+    model = _TinyAutoencoder()
+    snapshot = _snapshot_model(model)
+
+    try:
+        model.checkpoint_version = "unexpected-derived-metadata"
+        with pytest.raises(AssertionError):
+            _assert_model_unchanged(model, snapshot)
+    finally:
+        del model.checkpoint_version
+
+    _assert_model_unchanged(model, snapshot)
+    original_decoder_widths = model.mutable_config["decoder"]["widths"]
+    try:
+        model.mutable_config["decoder"]["widths"] = model.mutable_config["encoder"][
+            "widths"
+        ]
+        assert model.mutable_config == snapshot["metadata"]["mutable_config"]
+        with pytest.raises(AssertionError):
+            _assert_model_unchanged(model, snapshot)
+    finally:
+        model.mutable_config["decoder"]["widths"] = original_decoder_widths
+
+    _assert_model_unchanged(model, snapshot)
+
+
+def _legacy_factories_and_inputs():
+    return [
+        (
+            lambda: StandardAutoencoder(
+                k=2,
+                hidden_dims=[6, 4],
+                device="cpu",
+            ),
+            np.zeros((8, 3), dtype="float32"),
+        ),
+        (
+            lambda: OrthogonalAutoencoder(
+                k=2,
+                hidden_dims=[6, 4],
+                device="cpu",
+            ),
+            np.zeros((8, 3), dtype="float32"),
+        ),
+        (
+            lambda: LSTMAutoencoder(
+                k=2,
+                hidden=(5, 4),
+                device="cpu",
+            ),
+            np.zeros((8, 3, 2), dtype="float32"),
+        ),
+        (
+            lambda: CNNAutoencoder(k=2, device="cpu"),
+            np.zeros((8, 1, 4, 4), dtype="float32"),
+        ),
+        (
+            lambda: VisionTransformerAutoencoder(
+                k=2,
+                patch_size=2,
+                d_model=4,
+                depth_enc=1,
+                depth_dec=1,
+                heads=1,
+                device="cpu",
+            ),
+            np.zeros((8, 1, 4, 4), dtype="float32"),
+        ),
+        (
+            lambda: ConvLSTMAutoencoder(k=2, device="cpu"),
+            np.zeros((8, 2, 1, 4, 4), dtype="float32"),
+        ),
+        (
+            lambda: HybridConvLSTMTransformerAutoencoder(
+                k=2,
+                d_model=4,
+                n_heads=1,
+                n_layers=1,
+                efficient_attention=None,
+                device="cpu",
+            ),
+            np.zeros((8, 2, 1, 4, 4), dtype="float32"),
+        ),
+        (
+            lambda: VariationalAutoencoder(
+                k=2,
+                hidden_dims=[6, 4],
+                device="cpu",
+            ),
+            np.zeros((8, 3), dtype="float32"),
+        ),
+        (
+            lambda: SpatialTokenConvLSTMTransformerAutoencoder(
+                k=2,
+                spatial_pool_size=(1, 1),
+                d_model=4,
+                n_heads=1,
+                n_layers=1,
+                device="cpu",
+            ),
+            np.zeros((8, 2, 1, 4, 4), dtype="float32"),
+        ),
+    ]
+
+
+def _learning_rate_factories_and_inputs():
+    return [
+        (
+            lambda: StandardAutoencoder(
+                k=2,
+                hidden_dims=[4],
+                device="cpu",
+            ),
+            np.zeros((8, 3), dtype="float32"),
+        ),
+        (
+            lambda: OrthogonalAutoencoder(
+                k=2,
+                hidden_dims=[4],
+                lambda_W=0.0,
+                lambda_Z=0.0,
+                device="cpu",
+            ),
+            np.zeros((8, 3), dtype="float32"),
+        ),
+        (
+            lambda: VariationalAutoencoder(
+                k=2,
+                hidden_dims=[4],
+                beta=0.1,
+                device="cpu",
+            ),
+            np.zeros((8, 3), dtype="float32"),
+        ),
+    ]
+
+
+def _unbuilt_learning_rate_factories_and_inputs():
+    return [
+        (
+            lambda: _TinyAutoencoder(device="cpu"),
+            np.zeros((8, 3), dtype="float32"),
+        ),
+        (
+            lambda: OrthogonalAutoencoder(
+                k=2,
+                hidden_dims=[4],
+                lambda_W=0.0,
+                lambda_Z=0.0,
+                device="cpu",
+            ),
+            np.zeros((8, 3), dtype="float32"),
+        ),
+        (
+            lambda: VariationalAutoencoder(
+                k=2,
+                hidden_dims=[4],
+                beta=0.1,
+                device="cpu",
+            ),
+            np.zeros((8, 3), dtype="float32"),
+        ),
+    ]
+
+
+def _assert_unbuilt_training_state(model):
+    assert model.model is None
+    assert model._build_input_shape is None
+    assert model.is_fitted is False
+    for name in (
+        "optimizer",
+        "_optimizer",
+        "history",
+        "_history",
+        "training_history",
+        "best_model_state",
+        "_best_model_state",
+    ):
+        assert name not in model.__dict__
+
+
+@pytest.mark.parametrize(
+    "learning_rate",
+    [
+        True,
+        False,
+        -1.0,
+        float("nan"),
+        float("inf"),
+        -float("inf"),
+        "1e-3",
+        1e-3 + 0j,
+        [1e-3],
+        np.array([1e-3]),
+        None,
+    ],
+)
+@pytest.mark.parametrize(
+    ("factory", "X"),
+    _learning_rate_factories_and_inputs(),
+)
+def test_invalid_learning_rates_do_not_mutate_fit_state(
+    factory,
+    X,
+    learning_rate,
+):
+    model = factory()
+    model._build_input_shape = tuple(X.shape)
+    model.model = model._build_model(X.shape).to(model.device)
+    model.is_fitted = True
+    snapshot = _snapshot_model(model)
+
+    with pytest.raises(ValueError, match="learning_rate"):
+        model.fit(
+            X,
+            validation_split=0.25,
+            epochs=1,
+            batch_size=4,
+            learning_rate=learning_rate,
+            patience=2,
+            verbose=0,
+        )
+
+    _assert_model_unchanged(model, snapshot)
+
+
+@pytest.mark.parametrize(
+    "learning_rate",
+    [
+        True,
+        False,
+        -1.0,
+        float("nan"),
+        float("inf"),
+        -float("inf"),
+        "1e-3",
+        1e-3 + 0j,
+        np.array([1e-3]),
+    ],
+)
+@pytest.mark.parametrize(
+    ("factory", "X"),
+    _unbuilt_learning_rate_factories_and_inputs(),
+)
+def test_invalid_learning_rates_leave_unbuilt_models_pristine(
+    factory,
+    X,
+    learning_rate,
+):
+    model = factory()
+    _assert_unbuilt_training_state(model)
+    snapshot = _snapshot_model(model)
+
+    with pytest.raises(ValueError, match="learning_rate"):
+        model.fit(
+            X,
+            validation_split=0.25,
+            epochs=1,
+            batch_size=4,
+            learning_rate=learning_rate,
+            patience=2,
+            verbose=0,
+        )
+
+    _assert_model_unchanged(model, snapshot)
+    _assert_unbuilt_training_state(model)
+
+
+@pytest.mark.parametrize("learning_rate", [0.0, 1e-3])
+@pytest.mark.parametrize(
+    ("factory", "X"),
+    _learning_rate_factories_and_inputs(),
+)
+def test_valid_learning_rates_fit_all_training_paths(
+    factory,
+    X,
+    learning_rate,
+):
+    model = factory()
+
+    history = model.fit(
+        X,
+        validation_split=0.25,
+        epochs=1,
+        batch_size=4,
+        learning_rate=learning_rate,
+        patience=2,
+        verbose=0,
+    )
+
+    assert model.is_fitted
+    assert len(history["train_loss"]) == 1
+    assert len(history["val_loss"]) == 1
+
+
+@pytest.mark.parametrize(
+    ("factory", "X"),
+    _legacy_factories_and_inputs(),
+)
+def test_all_autoencoders_share_latent_input_validation(factory, X):
+    model = factory()
+    model._build_input_shape = tuple(X.shape)
+    model.model = model._build_model(X.shape).to(model.device)
+    model.is_fitted = True
+
+    single = model.decode(np.zeros(model.k, dtype="float32"), verbose=0)
+    batch = model.decode(np.zeros((3, model.k), dtype="float32"), verbose=0)
+    assert single.shape == (1, *X.shape[1:])
+    assert batch.shape == (3, *X.shape[1:])
+
+    invalid_shapes = [
+        np.zeros((2, 1, model.k), dtype="float32"),
+        np.zeros((2, model.k + 1), dtype="float32"),
+        np.zeros((2, 0), dtype="float32"),
+        np.zeros((0, model.k), dtype="float32"),
+    ]
+    for invalid in invalid_shapes:
+        with pytest.raises(ValueError, match="Z|latent"):
+            model.decode(invalid, verbose=0)
+
+    for invalid in [float("nan"), float("inf"), -float("inf")]:
+        latent = np.zeros((2, model.k), dtype="float32")
+        latent[0, 0] = invalid
+        with pytest.raises(ValueError, match="finite"):
+            model.decode(latent, verbose=0)
+
+    with pytest.raises(TypeError, match="real-valued"):
+        model.decode(np.zeros((2, model.k), dtype="complex64"), verbose=0)
+    with pytest.raises(TypeError, match="numeric"):
+        model.decode(np.full((2, model.k), "bad"), verbose=0)
+
+
+@pytest.mark.parametrize(
+    ("factory", "X"),
+    _legacy_factories_and_inputs(),
+)
+def test_all_autoencoders_reject_broadcastable_targets(factory, X):
+    model = factory()
+    target = X[(slice(None),) + (slice(0, 1),) * (X.ndim - 1)]
+
+    with pytest.raises(ValueError, match="Target shape"):
+        model.fit(
+            X,
+            y=target,
+            validation_split=0.25,
+            epochs=1,
+            batch_size=4,
+            patience=2,
+            verbose=0,
+        )
+
+
+@pytest.mark.parametrize(
+    ("factory", "X", "bad_X"),
+    [
+        (
+            lambda: StandardAutoencoder(k=2, hidden_dims=[4], device="cpu"),
+            np.zeros((4, 2, 5), dtype="float32"),
+            np.zeros((2, 10), dtype="float32"),
+        ),
+        (
+            lambda: OrthogonalAutoencoder(k=2, hidden_dims=[4], device="cpu"),
+            np.zeros((4, 2, 5), dtype="float32"),
+            np.zeros((2, 10), dtype="float32"),
+        ),
+        (
+            lambda: LSTMAutoencoder(k=2, hidden=(4, 3), device="cpu"),
+            np.zeros((4, 3, 2), dtype="float32"),
+            np.zeros((2, 4, 2), dtype="float32"),
+        ),
+        (
+            lambda: CNNAutoencoder(k=2, device="cpu"),
+            np.zeros((4, 1, 7, 9), dtype="float32"),
+            np.zeros((2, 1, 6, 10), dtype="float32"),
+        ),
+        (
+            lambda: VisionTransformerAutoencoder(
+                k=2,
+                patch_size=2,
+                d_model=4,
+                depth_enc=1,
+                depth_dec=1,
+                heads=1,
+                device="cpu",
+            ),
+            np.zeros((4, 1, 7, 9), dtype="float32"),
+            np.zeros((2, 1, 6, 10), dtype="float32"),
+        ),
+    ],
+)
+def test_predict_and_encode_reject_changed_sample_shapes(factory, X, bad_X):
+    model = factory()
+    model._build_input_shape = tuple(X.shape)
+    model.model = model._build_model(X.shape).to(model.device)
+    model.is_fitted = True
+
+    with pytest.raises(ValueError, match="Expected per-sample shape"):
+        model.predict(bad_X, verbose=0)
+    with pytest.raises(ValueError, match="Expected per-sample shape"):
+        model.encode(bad_X, verbose=0)
+
+
+@pytest.mark.parametrize("invalid", [np.nan, np.inf, -np.inf])
+def test_nonfinite_training_data_is_rejected(invalid):
+    X = np.zeros((8, 3), dtype="float32")
+    X[0, 0] = invalid
+    model = StandardAutoencoder(k=2, hidden_dims=[4], device="cpu")
+
+    with pytest.raises(ValueError, match="finite"):
+        model.fit(
+            X,
+            validation_split=0.25,
+            epochs=1,
+            batch_size=4,
+            patience=2,
+            verbose=0,
+        )
+    assert not model.is_fitted
+
+
+def test_nonfinite_custom_targets_are_rejected():
+    X = np.zeros((8, 3), dtype="float32")
+    y = X.copy()
+    y[0, 0] = np.nan
+    model = StandardAutoencoder(k=2, hidden_dims=[4], device="cpu")
+
+    with pytest.raises(ValueError, match="finite"):
+        model.fit(
+            X,
+            y=y,
+            validation_split=0.25,
+            epochs=1,
+            batch_size=4,
+            patience=2,
+            verbose=0,
+        )
+
+
+@pytest.mark.parametrize("operation", ["fit", "predict", "decode"])
+def test_float64_values_must_remain_finite_after_float32_cast(operation):
+    X = np.zeros((8, 3), dtype="float32")
+    huge = np.full((8, 3), 1e100, dtype="float64")
+    model = _TinyAutoencoder()
+
+    if operation != "fit":
+        model.fit(
+            X,
+            validation_split=0.25,
+            epochs=1,
+            batch_size=4,
+            learning_rate=0.0,
+            patience=2,
+            verbose=0,
+        )
+
+    with pytest.raises(ValueError, match="converted to float32"):
+        if operation == "fit":
+            model.fit(
+                huge,
+                validation_split=0.25,
+                epochs=1,
+                batch_size=4,
+                patience=2,
+                verbose=0,
+            )
+        elif operation == "predict":
+            model.predict(huge, verbose=0)
+        else:
+            model.decode(huge, verbose=0)
+
+
+def test_reconstruction_metrics_reject_invalid_custom_targets():
+    X = np.zeros((8, 3), dtype="float32")
+    model = _TinyAutoencoder()
+    model.fit(
+        X,
+        validation_split=0.25,
+        epochs=1,
+        batch_size=4,
+        learning_rate=0.0,
+        patience=2,
+        verbose=0,
+    )
+
+    with pytest.raises(ValueError, match="Target shape"):
+        model.reconstruction_error(X, y=X[:, :1])
+
+    nonfinite_target = X.copy()
+    nonfinite_target[0, 0] = np.nan
+    with pytest.raises(ValueError, match="finite"):
+        model.evaluate_reconstruction(X, y=nonfinite_target)
+
+
+@pytest.mark.parametrize(
+    "eps",
+    [-1.0, float("nan"), float("inf"), -float("inf"), True, "0.0", [0.0]],
+)
+@pytest.mark.parametrize(
+    "operation",
+    ["reconstruction_error", "evaluate_reconstruction", "evaluate"],
+)
+def test_model_metric_wrappers_validate_eps_before_prediction(operation, eps):
+    X = np.zeros((4, 3), dtype="float32")
+    model = _TinyAutoencoder()
+    model._build_input_shape = tuple(X.shape)
+    model.model = model._build_model(X.shape)
+    model.model.train()
+    model.is_fitted = True
+
+    with pytest.raises(ValueError, match="eps"):
+        getattr(model, operation)(X, eps=eps)
+
+    assert model.model.training
+
+
+class _NegatingAutoencoder(BaseDeepLearningModel):
+    def __init__(self):
+        super().__init__(device="cpu")
+
+    def _build_model(self, input_shape, **kwargs):
+        class _Model(nn.Module):
+            def __init__(self):
+                super().__init__()
+                self.anchor = nn.Parameter(torch.zeros(()))
+
+            def forward(self, x):
+                return -x + self.anchor * 0
+
+        return _Model()
+
+
+@pytest.mark.parametrize(
+    "operation",
+    ["reconstruction_error", "evaluate_reconstruction", "evaluate"],
+)
+def test_model_metric_wrappers_reject_arithmetic_overflow(operation):
+    X = np.full((4, 3), np.finfo(np.float32).max, dtype="float32")
+    model = _NegatingAutoencoder()
+    model._build_input_shape = tuple(X.shape)
+    model.model = model._build_model(X.shape)
+    model.is_fitted = True
+
+    with pytest.raises(FloatingPointError, match="Metric subtraction"):
+        getattr(model, operation)(X)
+
+
+class _NonfiniteInferenceAutoencoder(BaseDeepLearningModel):
+    def __init__(self):
+        super().__init__(device="cpu")
+
+    def _build_model(self, input_shape, **kwargs):
+        class _Model(nn.Module):
+            def __init__(self):
+                super().__init__()
+                self.anchor = nn.Parameter(torch.zeros(()))
+
+            def forward(self, x):
+                return torch.full_like(x, float("inf")) + self.anchor * 0
+
+            def encode_forward(self, x):
+                shape = (len(x), int(np.prod(x.shape[1:])))
+                return (
+                    torch.full(
+                        shape,
+                        float("inf"),
+                        dtype=x.dtype,
+                        device=x.device,
+                    )
+                    + self.anchor * 0
+                )
+
+            def decode_forward(self, z):
+                return torch.full_like(z, float("inf")) + self.anchor * 0
+
+        return _Model()
+
+
+@pytest.mark.parametrize("operation", ["predict", "encode", "decode"])
+def test_nonfinite_inference_results_are_rejected(operation):
+    X = np.zeros((4, 3), dtype="float32")
+    model = _NonfiniteInferenceAutoencoder()
+    model._build_input_shape = tuple(X.shape)
+    model.model = model._build_model(X.shape)
+    model.is_fitted = True
+
+    with pytest.raises(FloatingPointError, match="not finite"):
+        getattr(model, operation)(X, verbose=0)
+
+
+class _InfiniteBufferAutoencoder(BaseDeepLearningModel):
+    def __init__(self):
+        super().__init__(device="cpu")
+
+    def _build_model(self, input_shape, **kwargs):
+        class _Model(nn.Module):
+            def __init__(self):
+                super().__init__()
+                self.anchor = nn.Parameter(torch.zeros(()))
+                self.register_buffer("running_value", torch.zeros(()))
+
+            def forward(self, x):
+                self.running_value.fill_(float("inf"))
+                return torch.zeros_like(x) + self.anchor * 0
+
+        return _Model()
+
+
+def test_nonfinite_model_buffers_abort_training():
+    X = np.ones((8, 2), dtype="float32")
+    model = _InfiniteBufferAutoencoder()
+
+    with pytest.raises(FloatingPointError, match="Buffer"):
+        model.fit(
+            X,
+            validation_split=0.25,
+            epochs=1,
+            batch_size=4,
+            patience=2,
+            verbose=0,
+        )
+    assert not model.is_fitted
+
+
+class _ValidationOnlyInfiniteBufferAutoencoder(BaseDeepLearningModel):
+    def __init__(self):
+        super().__init__(device="cpu")
+
+    def _build_model(self, input_shape, **kwargs):
+        class _Model(nn.Module):
+            def __init__(self):
+                super().__init__()
+                self.anchor = nn.Parameter(torch.zeros(()))
+                self.register_buffer("running_value", torch.zeros(()))
+
+            def forward(self, x):
+                if not self.training:
+                    self.running_value.fill_(float("inf"))
+                return torch.zeros_like(x) + self.anchor * 0
+
+        return _Model()
+
+
+def test_validation_only_nonfinite_buffers_abort_training():
+    X = np.ones((8, 2), dtype="float32")
+    model = _ValidationOnlyInfiniteBufferAutoencoder()
+
+    with pytest.raises(FloatingPointError, match="Buffer"):
+        model.fit(
+            X,
+            validation_split=0.25,
+            epochs=1,
+            batch_size=4,
+            patience=2,
+            verbose=0,
+        )
+    assert not model.is_fitted
+
+
+def test_inference_only_nonfinite_buffers_are_rejected():
+    X = np.ones((4, 2), dtype="float32")
+    model = _ValidationOnlyInfiniteBufferAutoencoder()
+    model._build_input_shape = tuple(X.shape)
+    model.model = model._build_model(X.shape)
+    model.is_fitted = True
+
+    with pytest.raises(FloatingPointError, match="Buffer"):
+        model.predict(X, verbose=0)
+
+
+class _ValidationOnlyInfiniteParameterAutoencoder(BaseDeepLearningModel):
+    def __init__(self):
+        super().__init__(device="cpu")
+
+    def _build_model(self, input_shape, **kwargs):
+        class _Model(nn.Module):
+            def __init__(self):
+                super().__init__()
+                self.anchor = nn.Parameter(torch.zeros(()))
+
+            def forward(self, x):
+                if not self.training:
+                    self.anchor.fill_(float("inf"))
+                    return torch.zeros_like(x)
+                return torch.zeros_like(x) + self.anchor * 0
+
+        return _Model()
+
+
+def test_validation_only_nonfinite_parameters_abort_training():
+    X = np.ones((8, 2), dtype="float32")
+    model = _ValidationOnlyInfiniteParameterAutoencoder()
+
+    with pytest.raises(FloatingPointError, match="Parameter"):
+        model.fit(
+            X,
+            validation_split=0.25,
+            epochs=1,
+            batch_size=4,
+            patience=2,
+            verbose=0,
+        )
+    assert not model.is_fitted
+
+
+def test_inference_only_nonfinite_parameters_are_rejected():
+    X = np.ones((4, 2), dtype="float32")
+    model = _ValidationOnlyInfiniteParameterAutoencoder()
+    model._build_input_shape = tuple(X.shape)
+    model.model = model._build_model(X.shape)
+    model.is_fitted = True
+
+    with pytest.raises(FloatingPointError, match="Parameter"):
+        model.predict(X, verbose=0)
+
+
+class _NonTensorOutputAutoencoder(BaseDeepLearningModel):
+    def __init__(self):
+        super().__init__(device="cpu")
+
+    def _build_model(self, input_shape, **kwargs):
+        class _Model(nn.Module):
+            def __init__(self):
+                super().__init__()
+                self.anchor = nn.Parameter(torch.zeros(()))
+
+            def forward(self, x):
+                return (x + self.anchor * 0,)
+
+        return _Model()
+
+
+def test_non_tensor_model_outputs_are_rejected_clearly():
+    X = np.zeros((8, 3), dtype="float32")
+    model = _NonTensorOutputAutoencoder()
+
+    with pytest.raises(TypeError, match="Model output"):
+        model.fit(
+            X,
+            validation_split=0.25,
+            epochs=1,
+            batch_size=4,
+            patience=2,
+            verbose=0,
+        )
+    assert not model.is_fitted
+
+
+class _WrongShapeAutoencoder(BaseDeepLearningModel):
+    def __init__(self):
+        super().__init__(device="cpu")
+
+    def _build_model(self, input_shape, **kwargs):
+        class _Model(nn.Module):
+            def __init__(self):
+                super().__init__()
+                self.anchor = nn.Parameter(torch.zeros(()))
+
+            def forward(self, x):
+                return x[:, :1] + self.anchor * 0
+
+        return _Model()
+
+
+def test_model_outputs_cannot_broadcast_against_targets():
+    X = np.zeros((8, 3), dtype="float32")
+    model = _WrongShapeAutoencoder()
+
+    with pytest.raises(ValueError, match="output shape"):
+        model.fit(
+            X,
+            validation_split=0.25,
+            epochs=1,
+            batch_size=4,
+            patience=2,
+            verbose=0,
+        )
+    assert not model.is_fitted
+
+
+def test_nonfinite_loss_aborts_without_marking_model_fitted():
+    class _NaNLoss(nn.Module):
+        reduction = "mean"
+
+        def forward(self, output, target):
+            return output.sum() * torch.tensor(float("nan"))
+
+    X = np.ones((8, 2), dtype="float32")
+    model = _TinyAutoencoder()
+
+    with pytest.raises(FloatingPointError, match="not finite"):
+        model.fit(
+            X,
+            validation_split=0.25,
+            epochs=1,
+            batch_size=3,
+            criterion=_NaNLoss(),
+            patience=2,
+            verbose=0,
+        )
+    assert not model.is_fitted
+
+
+class _InfiniteGradientLoss(nn.Module):
+    reduction = "mean"
+
+    class _Operation(torch.autograd.Function):
+        @staticmethod
+        def forward(ctx, value):
+            return value.sum() * 0
+
+        @staticmethod
+        def backward(ctx, gradient):
+            return torch.full_like(gradient, float("inf"))
+
+    def forward(self, output, target):
+        return self._Operation.apply(output.sum())
+
+
+def test_nonfinite_gradients_abort_without_marking_model_fitted():
+    X = np.ones((8, 2), dtype="float32")
+    model = _TinyAutoencoder()
+
+    with pytest.raises(FloatingPointError, match="Gradient"):
+        model.fit(
+            X,
+            validation_split=0.25,
+            epochs=1,
+            batch_size=4,
+            criterion=_InfiniteGradientLoss(),
+            patience=2,
+            verbose=0,
+        )
+    assert not model.is_fitted
+
+
+class _InfiniteParameterOptimizer(torch.optim.Optimizer):
+    def __init__(self, parameters):
+        super().__init__(parameters, {})
+
+    @torch.no_grad()
+    def step(self, closure=None):
+        for group in self.param_groups:
+            for parameter in group["params"]:
+                parameter.fill_(float("inf"))
+
+
+def test_nonfinite_parameters_abort_without_marking_model_fitted():
+    X = np.ones((8, 2), dtype="float32")
+    model = _TinyAutoencoder()
+    model.model = model._build_model(X.shape)
+    optimizer = _InfiniteParameterOptimizer(model.model.parameters())
+
+    with pytest.raises(FloatingPointError, match="Parameter"):
+        model.fit(
+            X,
+            validation_split=0.25,
+            epochs=1,
+            batch_size=4,
+            optimizer=optimizer,
+            patience=2,
+            verbose=0,
+        )
+    assert not model.is_fitted
+
+
+def test_epoch_histories_weight_short_batches_by_sample_count():
+    X = np.array([[0.0], [1.0], [2.0], [3.0], [4.0], [5.0], [20.0]])
+    X = X.astype("float32")
+    validation_split = 0.4
+
+    np.random.seed(31)
+    indices = np.arange(len(X))
+    np.random.shuffle(indices)
+    split = int((1 - validation_split) * len(X))
+    train_indices = indices[:split]
+    validation_indices = indices[split:]
+    expected_train = float(np.mean(X[train_indices] ** 2))
+    expected_validation = float(np.mean(X[validation_indices] ** 2))
+
+    np.random.seed(31)
+    model = _TinyAutoencoder()
+    history = model.fit(
+        X,
+        validation_split=validation_split,
+        epochs=1,
+        batch_size=3,
+        learning_rate=0.0,
+        patience=2,
+        verbose=0,
+    )
+
+    assert history["train_loss"][0] == pytest.approx(expected_train)
+    assert history["val_loss"][0] == pytest.approx(expected_validation)
+
+
+def test_sum_reduced_histories_weight_short_batches_by_sample_count():
+    X = np.array([[0.0], [1.0], [2.0], [3.0], [4.0], [5.0], [20.0]])
+    X = X.astype("float32")
+    validation_split = 0.4
+
+    np.random.seed(43)
+    indices = np.arange(len(X))
+    np.random.shuffle(indices)
+    split = int((1 - validation_split) * len(X))
+    expected_train = float(np.mean(X[indices[:split]] ** 2))
+    expected_validation = float(np.mean(X[indices[split:]] ** 2))
+
+    np.random.seed(43)
+    model = _TinyAutoencoder()
+    history = model.fit(
+        X,
+        validation_split=validation_split,
+        epochs=1,
+        batch_size=3,
+        learning_rate=0.0,
+        criterion=nn.MSELoss(reduction="sum"),
+        patience=2,
+        verbose=0,
+    )
+
+    assert history["train_loss"][0] == pytest.approx(expected_train)
+    assert history["val_loss"][0] == pytest.approx(expected_validation)
+
+
+def test_orthogonal_histories_weight_short_batches_by_sample_count():
+    class _ZeroRegularizedModel(nn.Module):
+        def __init__(self):
+            super().__init__()
+            self.anchor = nn.Parameter(torch.zeros(()))
+
+        def forward(self, x):
+            return torch.zeros_like(x) + self.anchor * 0
+
+        def get_regularization_losses(self):
+            zero = self.anchor * 0
+            return zero, zero
+
+    X = np.array([[0.0], [1.0], [2.0], [3.0], [4.0], [5.0], [20.0]])
+    X = X.astype("float32")
+    validation_split = 0.4
+
+    np.random.seed(37)
+    indices = np.arange(len(X))
+    np.random.shuffle(indices)
+    split = int((1 - validation_split) * len(X))
+    expected_train = float(np.mean(X[indices[:split]] ** 2))
+    expected_validation = float(np.mean(X[indices[split:]] ** 2))
+
+    np.random.seed(37)
+    model = OrthogonalAutoencoder(
+        k=1,
+        hidden_dims=[2],
+        lambda_W=0.0,
+        lambda_Z=0.0,
+        device="cpu",
+    )
+    model.model = _ZeroRegularizedModel()
+    history = model.fit(
+        X,
+        validation_split=validation_split,
+        epochs=1,
+        batch_size=3,
+        learning_rate=0.0,
+        patience=2,
+        verbose=0,
+    )
+
+    assert history["train_loss"][0] == pytest.approx(expected_train)
+    assert history["val_loss"][0] == pytest.approx(expected_validation)
+
+
+def test_nonfinite_checkpoint_state_is_rejected_before_loading(tmp_path):
+    X = np.zeros((8, 3), dtype="float32")
+    model = _TinyAutoencoder()
+    model.fit(
+        X,
+        validation_split=0.25,
+        epochs=1,
+        batch_size=4,
+        learning_rate=0.0,
+        patience=2,
+        verbose=0,
+    )
+    original_anchor = model.model.anchor.detach().clone()
+
+    checkpoint_path = tmp_path / "tiny.pt"
+    model.save_pytorch_model(checkpoint_path)
+    checkpoint = torch.load(checkpoint_path, weights_only=False)
+    checkpoint["model_state_dict"]["anchor"] = torch.tensor(float("nan"))
+    torch.save(checkpoint, checkpoint_path)
+
+    with pytest.raises(FloatingPointError, match="Checkpoint state entry"):
+        model.load_pytorch_model(checkpoint_path, weights_only=False)
+    assert model.is_fitted
+    assert torch.equal(model.model.anchor.detach(), original_anchor)
+
+    with pytest.raises(FloatingPointError, match="Checkpoint state entry"):
+        _TinyAutoencoder.from_pytorch_model(
+            checkpoint_path,
+            weights_only=False,
+        )
+
+
+def test_late_checkpoint_shape_mismatch_is_transactional(tmp_path):
+    X = np.zeros((8, 3), dtype="float32")
+    model = _TinyAutoencoder()
+    model.fit(
+        X,
+        validation_split=0.25,
+        epochs=1,
+        batch_size=4,
+        learning_rate=0.0,
+        patience=2,
+        verbose=0,
+    )
+    model.model.eval()
+    snapshot = _snapshot_model(model)
+    checkpoint_path = tmp_path / "late-shape-mismatch.pt"
+    model.save_pytorch_model(checkpoint_path)
+    checkpoint = torch.load(checkpoint_path, weights_only=False)
+    checkpoint["model_state_dict"]["anchor"] = torch.ones(())
+    checkpoint["model_state_dict"]["tail"] = torch.ones(3)
+    torch.save(checkpoint, checkpoint_path)
+
+    with pytest.raises(RuntimeError, match="tail.*shape"):
+        model.load_pytorch_model(checkpoint_path, weights_only=False)
+
+    _assert_model_unchanged(model, snapshot)
+
+
+def test_checkpoint_cast_overflow_is_transactional(tmp_path):
+    X = np.zeros((8, 3), dtype="float32")
+    model = _TinyAutoencoder()
+    model.fit(
+        X,
+        validation_split=0.25,
+        epochs=1,
+        batch_size=4,
+        learning_rate=0.0,
+        patience=2,
+        verbose=0,
+    )
+    snapshot = _snapshot_model(model)
+    checkpoint_path = tmp_path / "cast-overflow.pt"
+    model.save_pytorch_model(checkpoint_path)
+    checkpoint = torch.load(checkpoint_path, weights_only=False)
+    checkpoint["model_state_dict"]["anchor"] = torch.tensor(
+        1e300,
+        dtype=torch.float64,
+    )
+    torch.save(checkpoint, checkpoint_path)
+
+    with pytest.raises(FloatingPointError, match="after conversion"):
+        model.load_pytorch_model(checkpoint_path, weights_only=False)
+
+    _assert_model_unchanged(model, snapshot)
+    with pytest.raises(FloatingPointError, match="after conversion"):
+        _TinyAutoencoder.from_pytorch_model(
+            checkpoint_path,
+            weights_only=False,
+        )
+
+
+def test_from_pytorch_model_rejects_late_shape_mismatch(tmp_path):
+    X = np.zeros((8, 3), dtype="float32")
+    model = _TinyAutoencoder()
+    model.fit(
+        X,
+        validation_split=0.25,
+        epochs=1,
+        batch_size=4,
+        learning_rate=0.0,
+        patience=2,
+        verbose=0,
+    )
+    checkpoint_path = tmp_path / "from-late-shape-mismatch.pt"
+    model.save_pytorch_model(checkpoint_path)
+    checkpoint = torch.load(checkpoint_path, weights_only=False)
+    checkpoint["model_state_dict"]["anchor"] = torch.ones(())
+    checkpoint["model_state_dict"]["tail"] = torch.ones(3)
+    torch.save(checkpoint, checkpoint_path)
+
+    with pytest.raises(RuntimeError, match="tail.*shape"):
+        _TinyAutoencoder.from_pytorch_model(
+            checkpoint_path,
+            weights_only=False,
+        )
+
+
+def test_failed_unbuilt_load_preserves_constructor_configuration(tmp_path):
+    X = np.zeros((8, 3), dtype="float32")
+    source = StandardAutoencoder(k=2, hidden_dims=[4], device="cpu")
+    source.fit(
+        X,
+        validation_split=0.25,
+        epochs=1,
+        batch_size=4,
+        learning_rate=0.0,
+        patience=2,
+        verbose=0,
+    )
+    checkpoint_path = tmp_path / "unbuilt-config-mismatch.pt"
+    source.save_pytorch_model(checkpoint_path)
+    checkpoint = torch.load(checkpoint_path, weights_only=False)
+    state_keys = list(checkpoint["model_state_dict"])
+    checkpoint["model_state_dict"][state_keys[0]] = (
+        checkpoint["model_state_dict"][state_keys[0]] + 1
+    )
+    later_value = checkpoint["model_state_dict"][state_keys[-1]]
+    checkpoint["model_state_dict"][state_keys[-1]] = torch.zeros(
+        later_value.numel() + 1,
+        dtype=later_value.dtype,
+    )
+    torch.save(checkpoint, checkpoint_path)
+
+    target = StandardAutoencoder(k=5, hidden_dims=[7], device="cpu")
+    snapshot = _snapshot_model(target)
+
+    with pytest.raises(RuntimeError, match="shape"):
+        target.load_pytorch_model(checkpoint_path, weights_only=False)
+
+    _assert_model_unchanged(target, snapshot)
+    assert target.k == 5
+    assert target.hidden_dims == [7]
+
+
+@pytest.mark.parametrize(
+    "factory",
+    [
+        lambda: StandardAutoencoder(k=0),
+        lambda: StandardAutoencoder(k=2, hidden_dims=[4, 0]),
+        lambda: OrthogonalAutoencoder(lambda_W=-1.0),
+        lambda: OrthogonalAutoencoder(lambda_Z=float("nan")),
+        lambda: LSTMAutoencoder(hidden=(4,)),
+        lambda: CNNAutoencoder(k=0),
+        lambda: VisionTransformerAutoencoder(patch_size=0),
+        lambda: VisionTransformerAutoencoder(d_model=1, heads=1),
+        lambda: VisionTransformerAutoencoder(d_model=2, heads=1),
+        lambda: VisionTransformerAutoencoder(
+            d_model=10,
+            heads=3,
+        ),
+        lambda: ConvLSTMAutoencoder(k=0),
+        lambda: HybridConvLSTMTransformerAutoencoder(
+            d_model=1,
+            n_heads=1,
+        ),
+        lambda: HybridConvLSTMTransformerAutoencoder(
+            d_model=2,
+            n_heads=1,
+        ),
+        lambda: HybridConvLSTMTransformerAutoencoder(efficient_attention="linera"),
+    ],
+)
+def test_existing_autoencoders_reject_invalid_constructor_values(factory):
+    with pytest.raises(ValueError):
+        factory()
+
+
+def test_linear_attention_accepts_long_sequences_without_quadratic_scores():
+    layer = LinearSelfAttention(d_model=16, num_heads=4)
+    values = torch.randn(2, 512, 16)
+
+    output = layer(values)
+
+    assert output.shape == values.shape
+    assert torch.isfinite(output).all()
+
+
+@pytest.mark.parametrize(
+    "kwargs",
+    [
+        {"d_model": 0, "num_heads": 1},
+        {"d_model": 4, "num_heads": 0},
+        {"d_model": 5, "num_heads": 2},
+        {"d_model": True, "num_heads": 1},
+    ],
+)
+def test_linear_attention_rejects_invalid_constructor_values(kwargs):
+    with pytest.raises(ValueError):
+        LinearSelfAttention(**kwargs)
+
+
+def test_linear_attention_rejects_incompatible_input_shape():
+    layer = LinearSelfAttention(d_model=4, num_heads=1)
+
+    with pytest.raises(ValueError, match="expects input shape"):
+        layer(torch.randn(2, 8, 3))
diff --git a/tests/deeplearning/test_global_state_isolation.py b/tests/deeplearning/test_global_state_isolation.py
new file mode 100644
index 0000000..3fd53d1
--- /dev/null
+++ b/tests/deeplearning/test_global_state_isolation.py
@@ -0,0 +1,121 @@
+"""Regression coverage for branch-test global-state isolation."""
+
+from pathlib import Path
+
+import numpy as np
+import pytest
+
+torch = pytest.importorskip("torch")
+
+
+class _FailSelectedTestCalls:
+    def pytest_runtest_call(self, item):
+        pytest.fail(
+            f"deliberate fixture-cleanup probe for {item.nodeid}",
+            pytrace=False,
+        )
+
+
+def _capture_state():
+    return {
+        "threads": torch.get_num_threads(),
+        "numpy": np.random.get_state(),
+        "torch": torch.random.get_rng_state().clone(),
+        "cuda": (
+            [state.clone() for state in torch.cuda.get_rng_state_all()]
+            if torch.cuda.is_available()
+            else None
+        ),
+    }
+
+
+def _restore_state(state):
+    torch.set_num_threads(state["threads"])
+    np.random.set_state(state["numpy"])
+    torch.random.set_rng_state(state["torch"])
+    if state["cuda"] is not None:
+        torch.cuda.set_rng_state_all(state["cuda"])
+
+
+def _assert_state_equal(actual, expected):
+    assert actual["threads"] == expected["threads"]
+    assert actual["numpy"][0] == expected["numpy"][0]
+    assert np.array_equal(actual["numpy"][1], expected["numpy"][1])
+    assert actual["numpy"][2:] == expected["numpy"][2:]
+    assert torch.equal(actual["torch"], expected["torch"])
+    if expected["cuda"] is None:
+        assert actual["cuda"] is None
+    else:
+        assert len(actual["cuda"]) == len(expected["cuda"])
+        for actual_state, expected_state in zip(
+            actual["cuda"],
+            expected["cuda"],
+        ):
+            assert torch.equal(actual_state, expected_state)
+
+
+def test_branch_test_fixtures_restore_threads_and_rng_state(tmp_path):
+    original = _capture_state()
+    tests_dir = Path(__file__).resolve().parent
+    try:
+        torch.set_num_threads(2)
+        np.random.seed(991)
+        torch.manual_seed(991)
+        baseline = _capture_state()
+        selected_nodes = [
+            (
+                "test_advanced_autoencoder_integration.py",
+                "test_advanced_autoencoders_are_publicly_exported",
+            ),
+            (
+                "test_autoencoder_hardening.py",
+                "test_nonfinite_custom_targets_are_rejected",
+            ),
+            (
+                "test_metrics.py",
+                "test_reconstruction_error_numpy_sample_mse_matches_manual",
+            ),
+            (
+                "test_spatial_token_autoencoder.py",
+                "test_spatial_token_rejects_non_sequence_input",
+            ),
+            (
+                "test_variational_autoencoder.py",
+                "test_vae_kl_has_known_analytic_values",
+            ),
+        ]
+        selected_tests = [
+            f"{tests_dir / module_name}::{test_name}"
+            for module_name, test_name in selected_nodes
+        ]
+
+        passing_result = pytest.main(
+            [
+                "-q",
+                "-p",
+                "no:cacheprovider",
+                "--basetemp",
+                str(tmp_path / "passing-inner-pytest"),
+                *selected_tests,
+            ]
+        )
+
+        assert passing_result == pytest.ExitCode.OK
+        _assert_state_equal(_capture_state(), baseline)
+
+        failing_result = pytest.main(
+            [
+                "-q",
+                "-p",
+                "no:cacheprovider",
+                "--basetemp",
+                str(tmp_path / "failing-inner-pytest"),
+                *selected_tests,
+            ],
+            plugins=[_FailSelectedTestCalls()],
+        )
+
+        assert failing_result == pytest.ExitCode.TESTS_FAILED
+        _assert_state_equal(_capture_state(), baseline)
+    finally:
+        _restore_state(original)
diff --git a/tests/deeplearning/test_metrics.py b/tests/deeplearning/test_metrics.py
index 20b591e..01c9e32 100644
--- a/tests/deeplearning/test_metrics.py
+++ b/tests/deeplearning/test_metrics.py
@@ -1,16 +1,32 @@
+"""Tests for public reconstruction metrics and losses."""
+
 import numpy as np
 import pytest
 
 torch = pytest.importorskip("torch")
 
-from bluemath_tk.deeplearning.metrics import (
+from bluemath_tk.deeplearning.metrics import (  # noqa: E402
     ReconstructionLoss,
     evaluate_reconstruction,
     reconstruction_error,
 )
 
 
-torch.set_num_threads(1)
+@pytest.fixture(autouse=True)
+def _preserve_global_state():
+    previous_threads = torch.get_num_threads()
+    numpy_state = np.random.get_state()
+    torch_state = torch.random.get_rng_state()
+    cuda_states = torch.cuda.get_rng_state_all() if torch.cuda.is_available() else None
+    try:
+        torch.set_num_threads(1)
+        yield
+    finally:
+        torch.set_num_threads(previous_threads)
+        np.random.set_state(numpy_state)
+        torch.random.set_rng_state(torch_state)
+        if cuda_states is not None:
+            torch.cuda.set_rng_state_all(cuda_states)
 
 
 def test_reconstruction_error_numpy_sample_mse_matches_manual():
@@ -99,6 +115,18 @@ def test_evaluate_reconstruction_returns_summary_statistics():
     assert summary["min"] <= summary["mean"] <= summary["max"]
 
 
+def test_evaluate_reconstruction_uses_stable_float64_summary_mean(monkeypatch):
+    y_true = np.zeros((3, 1), dtype=np.float64)
+    y_pred = np.full((3, 1), 1e308, dtype=np.float64)
+
+    # Isolate mean accumulation from NumPy's independent extreme-value std overflow.
+    monkeypatch.setattr(np, "std", lambda _values: np.float64(0.0))
+    summary = evaluate_reconstruction(y_true, y_pred, metric="mae")
+
+    assert summary["n_samples"] == 3
+    assert summary["mean"] == pytest.approx(1e308)
+
+
 def test_reconstruction_error_torch_mean_is_differentiable():
     """Tensor inputs should keep gradients so the metric can be used as a loss."""
     y_true = torch.zeros((4, 3), dtype=torch.float32)
@@ -164,3 +192,626 @@ def test_reconstruction_metrics_validate_inputs():
 
     with pytest.raises(ValueError, match="reduction"):
         reconstruction_error(y_true, y_true, reduction="invalid")
+
+
+@pytest.mark.parametrize("dtype", [torch.float32, torch.float64])
+@pytest.mark.parametrize("reduction", ["none", "sample", "mean", "sum"])
+def test_exact_zero_rmse_has_finite_zero_gradient(reduction, dtype):
+    y_true = torch.zeros((3, 2, 4), dtype=dtype)
+    y_pred = torch.zeros((3, 2, 4), dtype=dtype, requires_grad=True)
+
+    loss = reconstruction_error(
+        y_true,
+        y_pred,
+        metric="rmse",
+        reduction=reduction,
+        eps=0.0,
+    )
+    loss.sum().backward()
+
+    assert torch.count_nonzero(loss) == 0
+    assert torch.isfinite(y_pred.grad).all()
+    assert torch.count_nonzero(y_pred.grad) == 0
+
+
+def test_reconstruction_loss_exact_zero_rmse_has_finite_zero_gradient():
+    y_true = torch.zeros((3, 5), dtype=torch.float32)
+    y_pred = torch.zeros((3, 5), dtype=torch.float32, requires_grad=True)
+    criterion = ReconstructionLoss(metric="rmse", reduction="mean", eps=0.0)
+
+    loss = criterion(y_pred, y_true)
+    loss.backward()
+
+    assert loss.item() == 0.0
+    assert torch.isfinite(y_pred.grad).all()
+    assert torch.count_nonzero(y_pred.grad) == 0
+
+
+def test_rmse_values_and_gradients_match_analytic_result_away_from_zero():
+    y_true = torch.zeros((2, 2, 2), dtype=torch.float64)
+    values = torch.tensor(
+        [[[1.0, 2.0], [3.0, 4.0]], [[2.0, 3.0], [4.0, 5.0]]],
+        dtype=torch.float64,
+    )
+    y_pred = values.clone().requires_grad_(True)
+    reference_pred = values.clone().requires_grad_(True)
+
+    actual = reconstruction_error(
+        y_true,
+        y_pred,
+        metric="rmse",
+        reduction="mean",
+        eps=0.0,
+    )
+    expected = torch.sqrt(reference_pred.pow(2).mean(dim=(1, 2))).mean()
+    actual.backward()
+    expected.backward()
+
+    assert torch.allclose(actual, expected)
+    assert torch.allclose(y_pred.grad, reference_pred.grad)
+    assert torch.isfinite(y_pred.grad).all()
+
+
+def _assert_constant_float32_rmse(residual, reduction):
+    y_true = torch.zeros((2, 2), dtype=torch.float32)
+    y_pred = torch.full(
+        (2, 2),
+        residual,
+        dtype=torch.float32,
+        requires_grad=True,
+    )
+
+    actual = reconstruction_error(
+        y_true,
+        y_pred,
+        metric="rmse",
+        reduction=reduction,
+        eps=0.0,
+    )
+    actual.sum().backward()
+
+    stored_residual = float(y_pred.detach()[0, 0])
+    expected_value = stored_residual * 2 if reduction == "sum" else stored_residual
+    expected = torch.full_like(actual, expected_value)
+    gradient_scale = {
+        "none": 1.0,
+        "sample": 0.5,
+        "mean": 0.25,
+        "sum": 0.5,
+    }[reduction]
+    if stored_residual == 0:
+        gradient_scale = 0.0
+    expected_gradient = torch.full_like(y_pred, gradient_scale)
+    numpy_result = reconstruction_error(
+        np.zeros((2, 2), dtype=np.float32),
+        np.full((2, 2), residual, dtype=np.float32),
+        metric="rmse",
+        reduction=reduction,
+        eps=0.0,
+    )
+
+    assert actual.dtype == torch.float32
+    assert actual.device == y_pred.device
+    assert torch.isfinite(actual).all()
+    torch.testing.assert_close(actual, expected, rtol=5e-6, atol=0.0)
+    np.testing.assert_allclose(
+        actual.detach().cpu().numpy(),
+        np.asarray(numpy_result),
+        rtol=5e-6,
+        atol=0.0,
+    )
+    assert y_pred.grad is not None
+    assert torch.isfinite(y_pred.grad).all()
+    torch.testing.assert_close(
+        y_pred.grad,
+        expected_gradient,
+        rtol=5e-6,
+        atol=0.0,
+    )
+
+
+@pytest.mark.parametrize(
+    "residual",
+    [
+        pytest.param(0.0, id="zero"),
+        pytest.param(1e-25, id="small"),
+        pytest.param(3.0, id="moderate"),
+    ],
+)
+@pytest.mark.parametrize("reduction", ["none", "sample", "mean", "sum"])
+def test_float32_rmse_multisample_reductions_are_stable(residual, reduction):
+    _assert_constant_float32_rmse(residual, reduction)
+
+
+@pytest.mark.parametrize("reduction", ["none", "sample", "mean"])
+def test_float32_extreme_rmse_multisample_mean_is_stable(reduction):
+    _assert_constant_float32_rmse(3e38, reduction)
+
+
+@pytest.mark.parametrize(
+    ("torch_dtype", "numpy_dtype", "metric", "residual"),
+    [
+        pytest.param(
+            torch.float32,
+            np.float32,
+            "mae",
+            3e38,
+            id="float32-mae",
+        ),
+        pytest.param(
+            torch.float32,
+            np.float32,
+            "rmse",
+            3e38,
+            id="float32-rmse",
+        ),
+        pytest.param(
+            torch.float32,
+            np.float32,
+            "mse",
+            1.8e19,
+            id="float32-mse",
+        ),
+        pytest.param(
+            torch.float64,
+            np.float64,
+            "mae",
+            1e308,
+            id="float64-mae",
+        ),
+        pytest.param(
+            torch.float64,
+            np.float64,
+            "rmse",
+            1e308,
+            id="float64-rmse",
+        ),
+        pytest.param(
+            torch.float64,
+            np.float64,
+            "mse",
+            1.3e154,
+            id="float64-mse",
+        ),
+    ],
+)
+def test_multisample_extreme_metric_means_match_numpy_and_gradients(
+    torch_dtype,
+    numpy_dtype,
+    metric,
+    residual,
+):
+    y_true = torch.zeros((2, 2), dtype=torch_dtype)
+    y_pred = torch.full(
+        (2, 2),
+        residual,
+        dtype=torch_dtype,
+        requires_grad=True,
+    )
+    numpy_true = np.zeros((2, 2), dtype=numpy_dtype)
+    numpy_pred = np.full((2, 2), residual, dtype=numpy_dtype)
+
+    sample_result = reconstruction_error(
+        y_true,
+        y_pred,
+        metric=metric,
+        reduction="sample",
+    )
+    mean_result = reconstruction_error(
+        y_true,
+        y_pred,
+        metric=metric,
+        reduction="mean",
+    )
+    numpy_sample = reconstruction_error(
+        numpy_true,
+        numpy_pred,
+        metric=metric,
+        reduction="sample",
+    )
+    numpy_mean = reconstruction_error(
+        numpy_true,
+        numpy_pred,
+        metric=metric,
+        reduction="mean",
+    )
+    mean_result.backward()
+
+    stored_residual = float(y_pred.detach()[0, 0])
+    if metric == "mse":
+        expected_value = stored_residual * stored_residual
+        expected_gradient = stored_residual / 2
+    else:
+        expected_value = stored_residual
+        expected_gradient = 0.25
+    rtol = 5e-6 if torch_dtype == torch.float32 else 5e-15
+
+    assert sample_result.shape == (2,)
+    assert sample_result.dtype == torch_dtype
+    assert mean_result.shape == ()
+    assert mean_result.dtype == torch_dtype
+    assert mean_result.device == y_pred.device
+    assert numpy_sample.shape == (2,)
+    assert numpy_sample.dtype == numpy_dtype
+    assert torch.isfinite(sample_result).all()
+    assert torch.isfinite(mean_result)
+    torch.testing.assert_close(
+        sample_result,
+        torch.full_like(sample_result, expected_value),
+        rtol=rtol,
+        atol=0.0,
+    )
+    torch.testing.assert_close(
+        mean_result,
+        torch.tensor(expected_value, dtype=torch_dtype),
+        rtol=rtol,
+        atol=0.0,
+    )
+    np.testing.assert_allclose(
+        numpy_sample,
+        np.full((2,), expected_value, dtype=numpy_dtype),
+        rtol=rtol,
+        atol=0.0,
+    )
+    np.testing.assert_allclose(
+        numpy_mean,
+        expected_value,
+        rtol=rtol,
+        atol=0.0,
+    )
+    assert y_pred.grad is not None
+    assert torch.isfinite(y_pred.grad).all()
+    torch.testing.assert_close(
+        y_pred.grad,
+        torch.full_like(y_pred, expected_gradient),
+        rtol=rtol,
+        atol=0.0,
+    )
+
+
+@pytest.mark.parametrize(
+    ("torch_dtype", "numpy_dtype", "residual"),
+    [
+        pytest.param(torch.float32, np.float32, 2e19, id="float32"),
+        pytest.param(torch.float64, np.float64, 1.5e154, id="float64"),
+    ],
+)
+def test_mse_grouping_avoids_nonrepresentable_elementwise_square(
+    torch_dtype,
+    numpy_dtype,
+    residual,
+):
+    y_true = torch.zeros((2, 2), dtype=torch_dtype)
+    y_pred = torch.tensor(
+        [[residual, 0.0], [residual, 0.0]],
+        dtype=torch_dtype,
+        requires_grad=True,
+    )
+    numpy_true = np.zeros((2, 2), dtype=numpy_dtype)
+    numpy_pred = np.array(
+        [[residual, 0.0], [residual, 0.0]],
+        dtype=numpy_dtype,
+    )
+
+    assert torch.isinf(y_pred.detach().square()).any()
+    with np.errstate(over="ignore"):
+        assert np.isinf(numpy_pred**2).any()
+
+    sample_result = reconstruction_error(
+        y_true,
+        y_pred,
+        metric="mse",
+        reduction="sample",
+    )
+    mean_result = reconstruction_error(
+        y_true,
+        y_pred,
+        metric="mse",
+        reduction="mean",
+    )
+    numpy_sample = reconstruction_error(
+        numpy_true,
+        numpy_pred,
+        metric="mse",
+        reduction="sample",
+    )
+    numpy_mean = reconstruction_error(
+        numpy_true,
+        numpy_pred,
+        metric="mse",
+        reduction="mean",
+    )
+    mean_result.backward()
+
+    stored_residual = float(y_pred.detach()[0, 0])
+    expected_value = stored_residual * (stored_residual / 2)
+    expected_gradient = y_pred.detach() / 2
+    rtol = 5e-6 if torch_dtype == torch.float32 else 5e-15
+
+    assert sample_result.dtype == torch_dtype
+    assert mean_result.dtype == torch_dtype
+    assert numpy_sample.dtype == numpy_dtype
+    assert torch.isfinite(sample_result).all()
+    assert torch.isfinite(mean_result)
+    torch.testing.assert_close(
+        sample_result,
+        torch.full_like(sample_result, expected_value),
+        rtol=rtol,
+        atol=0.0,
+    )
+    torch.testing.assert_close(
+        mean_result,
+        torch.tensor(expected_value, dtype=torch_dtype),
+        rtol=rtol,
+        atol=0.0,
+    )
+    np.testing.assert_allclose(
+        numpy_sample,
+        np.full((2,), expected_value, dtype=numpy_dtype),
+        rtol=rtol,
+        atol=0.0,
+    )
+    np.testing.assert_allclose(
+        numpy_mean,
+        expected_value,
+        rtol=rtol,
+        atol=0.0,
+    )
+    assert y_pred.grad is not None
+    assert torch.isfinite(y_pred.grad).all()
+    torch.testing.assert_close(
+        y_pred.grad,
+        expected_gradient,
+        rtol=rtol,
+        atol=0.0,
+    )
+
+
+@pytest.mark.parametrize("metric", ["mse", "mae"])
+@pytest.mark.parametrize("dtype", [torch.float32, torch.float64])
+def test_exact_zero_scaled_mean_has_zero_gradient(metric, dtype):
+    y_true = torch.zeros((2, 2), dtype=dtype)
+    y_pred = torch.zeros((2, 2), dtype=dtype, requires_grad=True)
+
+    result = reconstruction_error(
+        y_true,
+        y_pred,
+        metric=metric,
+        reduction="mean",
+    )
+    result.backward()
+
+    assert result.dtype == dtype
+    assert result.item() == 0.0
+    assert torch.isfinite(y_pred.grad).all()
+    assert torch.count_nonzero(y_pred.grad) == 0
+
+
+@pytest.mark.parametrize(
+    "residual",
+    [
+        pytest.param(1e-25, id="small"),
+        pytest.param(3e38, id="large"),
+    ],
+)
+def test_reconstruction_loss_rmse_preserves_extreme_float32_gradients(residual):
+    y_true = torch.zeros((2, 2), dtype=torch.float32)
+    y_pred = torch.full(
+        (2, 2),
+        residual,
+        dtype=torch.float32,
+        requires_grad=True,
+    )
+    criterion = ReconstructionLoss(metric="rmse", reduction="mean", eps=0.0)
+
+    loss = criterion(y_pred, y_true)
+    loss.backward()
+
+    torch.testing.assert_close(
+        loss,
+        torch.tensor(residual, dtype=torch.float32),
+        rtol=5e-6,
+        atol=0.0,
+    )
+    assert torch.isfinite(y_pred.grad).all()
+    torch.testing.assert_close(
+        y_pred.grad,
+        torch.full_like(y_pred, 0.25),
+        rtol=5e-6,
+        atol=0.0,
+    )
+
+
+def test_scaled_rmse_preserves_eps_inside_the_square_root():
+    y_true = torch.zeros((1, 2), dtype=torch.float32)
+    y_pred = torch.tensor([[3.0, 4.0]], requires_grad=True)
+    eps = 2.25
+
+    loss = reconstruction_error(
+        y_true,
+        y_pred,
+        metric="rmse",
+        reduction="sample",
+        eps=eps,
+    )
+    loss.sum().backward()
+
+    expected_value = np.sqrt((3.0**2 + 4.0**2) / 2 + eps)
+    expected_gradient = torch.tensor(
+        [[3.0 / (2 * expected_value), 4.0 / (2 * expected_value)]],
+        dtype=torch.float32,
+    )
+    torch.testing.assert_close(
+        loss,
+        torch.tensor([expected_value], dtype=torch.float32),
+    )
+    torch.testing.assert_close(y_pred.grad, expected_gradient)
+
+
+@pytest.mark.parametrize("kind", ["numpy", "torch"])
+def test_rmse_sum_rejects_a_nonrepresentable_float32_final_result(kind):
+    if kind == "torch":
+        y_true = torch.zeros((2, 1), dtype=torch.float32)
+        y_pred = torch.full((2, 1), 3e38, dtype=torch.float32)
+    else:
+        y_true = np.zeros((2, 1), dtype=np.float32)
+        y_pred = np.full((2, 1), 3e38, dtype=np.float32)
+
+    with pytest.raises(FloatingPointError, match="Metric reduction"):
+        reconstruction_error(
+            y_true,
+            y_pred,
+            metric="rmse",
+            reduction="sum",
+            eps=0.0,
+        )
+
+
+@pytest.mark.parametrize("metric", ["mse", "mae", "rmse"])
+@pytest.mark.parametrize(
+    "kind",
+    ["numpy", "torch", "numpy-target", "numpy-prediction"],
+)
+def test_metrics_support_valid_numpy_torch_and_mixed_inputs(metric, kind):
+    numpy_true = np.zeros((3, 2), dtype="float32")
+    numpy_pred = np.ones((3, 2), dtype="float32")
+    if kind == "numpy":
+        y_true, y_pred = numpy_true, numpy_pred
+    elif kind == "torch":
+        y_true = torch.from_numpy(numpy_true)
+        y_pred = torch.from_numpy(numpy_pred)
+    elif kind == "numpy-target":
+        y_true, y_pred = numpy_true, torch.from_numpy(numpy_pred)
+    else:
+        y_true, y_pred = torch.from_numpy(numpy_true), numpy_pred
+
+    result = reconstruction_error(y_true, y_pred, metric=metric, reduction="mean")
+    summary = evaluate_reconstruction(y_true, y_pred, metric=metric)
+
+    if torch.is_tensor(result):
+        assert result.item() == pytest.approx(1.0)
+    else:
+        assert result == pytest.approx(1.0)
+    assert summary["mean"] == pytest.approx(1.0)
+
+
+@pytest.mark.parametrize("invalid", [float("nan"), float("inf"), -float("inf")])
+@pytest.mark.parametrize("kind", ["numpy", "torch", "mixed"])
+def test_metrics_reject_nonfinite_inputs(invalid, kind):
+    y_true = np.zeros((2, 3), dtype="float32")
+    y_pred = np.zeros((2, 3), dtype="float32")
+    y_pred[0, 0] = invalid
+    if kind == "torch":
+        y_true = torch.from_numpy(y_true)
+        y_pred = torch.from_numpy(y_pred)
+    elif kind == "mixed":
+        y_pred = torch.from_numpy(y_pred)
+
+    with pytest.raises(ValueError, match="finite"):
+        reconstruction_error(y_true, y_pred)
+    with pytest.raises(ValueError, match="finite"):
+        evaluate_reconstruction(y_true, y_pred)
+
+
+@pytest.mark.parametrize(
+    "values",
+    [
+        np.zeros((2, 3), dtype="complex64"),
+        torch.zeros((2, 3), dtype=torch.complex64),
+    ],
+)
+def test_metrics_reject_complex_inputs(values):
+    zeros = np.zeros((2, 3), dtype="float32")
+
+    with pytest.raises(TypeError, match="real-valued"):
+        reconstruction_error(zeros, values)
+
+
+@pytest.mark.parametrize(
+    "values",
+    [
+        np.full((2, 3), "bad"),
+        np.ones((2, 3), dtype=bool),
+        torch.ones((2, 3), dtype=torch.bool),
+        [[0.0, 0.0], [0.0, 0.0]],
+    ],
+)
+def test_metrics_reject_unsupported_or_non_numeric_inputs(values):
+    zeros = np.zeros((2, 3), dtype="float32")
+
+    with pytest.raises(TypeError, match="numeric|NumPy array"):
+        reconstruction_error(zeros, values)
+
+
+@pytest.mark.parametrize(
+    "shape",
+    [(), (0,), (0, 3), (2, 0)],
+)
+def test_metrics_reject_empty_or_scalar_inputs(shape):
+    values = np.empty(shape, dtype="float32")
+
+    with pytest.raises(ValueError, match="sample dimension|non-empty"):
+        reconstruction_error(values, values)
+
+
+@pytest.mark.parametrize(
+    "eps",
+    [
+        -1.0,
+        float("nan"),
+        float("inf"),
+        -float("inf"),
+        True,
+        False,
+        "0.0",
+        1 + 0j,
+        [0.0],
+        np.array([0.0]),
+    ],
+)
+def test_all_public_metrics_reject_invalid_eps(eps):
+    values = np.zeros((2, 3), dtype="float32")
+
+    with pytest.raises(ValueError, match="eps"):
+        reconstruction_error(values, values, metric="rmse", eps=eps)
+    with pytest.raises(ValueError, match="eps"):
+        evaluate_reconstruction(values, values, metric="rmse", eps=eps)
+    with pytest.raises(ValueError, match="eps"):
+        ReconstructionLoss(metric="rmse", eps=eps)
+
+
+@pytest.mark.parametrize("kind", ["numpy", "torch", "mixed-cast"])
+def test_metrics_reject_nonfinite_arithmetic_overflow(kind):
+    if kind == "numpy":
+        limit = np.finfo(np.float64).max
+        y_true = np.full((2, 2), limit, dtype="float64")
+        y_pred = np.full((2, 2), -limit, dtype="float64")
+    elif kind == "torch":
+        limit = torch.finfo(torch.float32).max
+        y_true = torch.full((2, 2), limit, dtype=torch.float32)
+        y_pred = torch.full((2, 2), -limit, dtype=torch.float32)
+    else:
+        y_true = np.full((2, 2), 1e300, dtype="float64")
+        y_pred = torch.zeros((2, 2), dtype=torch.float32)
+
+    with pytest.raises(FloatingPointError, match="finite"):
+        reconstruction_error(y_true, y_pred, metric="mse")
+
+
+def test_evaluate_reconstruction_rejects_summary_overflow():
+    limit = np.finfo(np.float64).max
+    y_true = np.zeros((2, 1), dtype="float64")
+    y_pred = np.array([[limit], [0.0]], dtype="float64")
+
+    with pytest.raises(FloatingPointError, match="summary"):
+        evaluate_reconstruction(y_true, y_pred, metric="mae")
+
+
+def test_reconstruction_loss_validates_forward_inputs():
+    criterion = ReconstructionLoss(metric="mse")
+    y_true = torch.zeros((2, 3), dtype=torch.float32)
+    y_pred = torch.zeros((2, 3), dtype=torch.float32)
+    y_pred[0, 0] = float("nan")
+
+    with pytest.raises(ValueError, match="finite"):
+        criterion(y_pred, y_true)
diff --git a/tests/deeplearning/test_spatial_token_autoencoder.py b/tests/deeplearning/test_spatial_token_autoencoder.py
new file mode 100644
index 0000000..0abf84e
--- /dev/null
+++ b/tests/deeplearning/test_spatial_token_autoencoder.py
@@ -0,0 +1,306 @@
+"""Tests for the spatial-token spatiotemporal autoencoder."""
+
+import numpy as np
+import pytest
+import torch
+
+from bluemath_tk.deeplearning.autoencoders import (
+    SpatialTokenConvLSTMTransformerAutoencoder,
+)
+from bluemath_tk.deeplearning.spatiotemporal_autoencoders import (
+    _FactorizedSpatiotemporalBlock,
+)
+
+
+@pytest.fixture(autouse=True)
+def _set_seed():
+    previous_threads = torch.get_num_threads()
+    numpy_state = np.random.get_state()
+    torch_state = torch.random.get_rng_state()
+    cuda_states = torch.cuda.get_rng_state_all() if torch.cuda.is_available() else None
+    try:
+        np.random.seed(503)
+        torch.manual_seed(503)
+        torch.set_num_threads(1)
+        yield
+    finally:
+        torch.set_num_threads(previous_threads)
+        np.random.set_state(numpy_state)
+        torch.random.set_rng_state(torch_state)
+        if cuda_states is not None:
+            torch.cuda.set_rng_state_all(cuda_states)
+
+
+def _model():
+    return SpatialTokenConvLSTMTransformerAutoencoder(
+        k=4,
+        spatial_pool_size=(2, 2),
+        d_model=8,
+        n_heads=2,
+        n_layers=1,
+        device="cpu",
+    )
+
+
+def _fit_model(X):
+    model = _model()
+    model.fit(
+        X,
+        validation_split=0.25,
+        epochs=1,
+        batch_size=4,
+        learning_rate=1e-3,
+        patience=2,
+        verbose=0,
+    )
+    return model
+
+
+def test_spatial_token_autoencoder_reconstructs_complete_sequence():
+    X = np.random.randn(12, 3, 1, 8, 10).astype("float32")
+    model = _fit_model(X)
+
+    latent = model.encode(X, batch_size=4, verbose=0)
+    prediction = model.predict(X, batch_size=4, verbose=0)
+    decoded = model.decode(latent, batch_size=4, verbose=0)
+
+    assert latent.shape == (12, 4)
+    assert prediction.shape == X.shape
+    assert decoded.shape == X.shape
+    assert np.allclose(prediction, decoded, rtol=1e-5, atol=1e-6)
+    assert np.isfinite(prediction).all()
+
+
+def test_spatial_token_single_vector_decoding():
+    X = np.random.randn(12, 3, 1, 7, 9).astype("float32")
+    model = _fit_model(X)
+    latent = model.encode(X, verbose=0)
+
+    decoded = model.decode(latent[0], verbose=0)
+
+    assert decoded.shape == (1, 3, 1, 7, 9)
+
+
+def test_spatial_token_fixed_sample_shape_validation():
+    X = np.random.randn(12, 3, 1, 8, 8).astype("float32")
+    model = _fit_model(X)
+
+    with pytest.raises(ValueError, match="Expected per-sample shape"):
+        model.predict(
+            np.random.randn(4, 4, 1, 8, 8).astype("float32"),
+            verbose=0,
+        )
+    with pytest.raises(ValueError, match="Expected per-sample shape"):
+        model.encode(
+            np.random.randn(4, 3, 2, 8, 8).astype("float32"),
+            verbose=0,
+        )
+
+
+def test_spatial_token_rejects_non_sequence_input():
+    model = _model()
+    X = np.random.randn(12, 1, 8, 8).astype("float32")
+
+    with pytest.raises(ValueError, match="5D input"):
+        model.fit(
+            X,
+            validation_split=0.25,
+            epochs=1,
+            batch_size=4,
+            patience=2,
+            verbose=0,
+        )
+
+
+def test_spatial_token_rejects_non_sequence_targets():
+    X = np.random.randn(12, 3, 1, 8, 8).astype("float32")
+    for target in (X[:, -1], X[:, :1]):
+        with pytest.raises(ValueError, match="Target shape"):
+            _model().fit(
+                X,
+                y=target,
+                validation_split=0.25,
+                epochs=1,
+                batch_size=4,
+                patience=2,
+                verbose=0,
+            )
+
+
+def test_spatial_token_rejects_pool_larger_than_encoded_grid():
+    X = np.random.randn(12, 3, 1, 6, 6).astype("float32")
+    model = SpatialTokenConvLSTMTransformerAutoencoder(
+        k=4,
+        spatial_pool_size=(3, 3),
+        d_model=8,
+        n_heads=2,
+        n_layers=1,
+        device="cpu",
+    )
+
+    with pytest.raises(ValueError, match="spatial_pool_size"):
+        model.fit(
+            X,
+            validation_split=0.25,
+            epochs=1,
+            batch_size=4,
+            patience=2,
+            verbose=0,
+        )
+
+
+def test_spatial_token_singleton_tiny_grid_is_safe_and_input_sensitive():
+    outer = SpatialTokenConvLSTMTransformerAutoencoder(
+        k=2,
+        spatial_pool_size=(1, 1),
+        d_model=4,
+        n_heads=1,
+        n_layers=1,
+        device="cpu",
+    )
+    inner = outer._build_model((2, 1, 1, 1, 1)).eval()
+    inputs = torch.tensor(
+        [
+            [[[[0.0]]]],
+            [[[[1.0]]]],
+        ]
+    )
+    with torch.no_grad():
+        latent = inner.encode_forward(inputs)
+        reconstruction = inner(inputs)
+
+    assert reconstruction.shape == (2, 1, 1, 1, 1)
+    assert torch.isfinite(reconstruction).all()
+    assert not torch.allclose(latent[0], latent[1], atol=1e-6, rtol=1e-5)
+
+
+@pytest.mark.parametrize(
+    ("kwargs", "message"),
+    [
+        ({"k": 0}, "k"),
+        ({"spatial_pool_size": (0, 2)}, "spatial_pool_size"),
+        ({"spatial_pool_size": [2, 2]}, "spatial_pool_size"),
+        ({"d_model": 0}, "d_model"),
+        ({"d_model": 1}, "d_model"),
+        ({"d_model": 2}, "d_model"),
+        ({"n_heads": 0}, "n_heads"),
+        ({"d_model": 10, "n_heads": 3}, "divisible"),
+        ({"n_layers": 0}, "n_layers"),
+    ],
+)
+def test_spatial_token_rejects_invalid_constructor_arguments(kwargs, message):
+    with pytest.raises(ValueError, match=message):
+        SpatialTokenConvLSTMTransformerAutoencoder(**kwargs)
+
+
+def test_spatial_token_gradients_reach_space_time_and_latent_paths():
+    X = torch.randn(4, 3, 1, 6, 6)
+    model = _model()._build_model(tuple(X.shape))
+
+    reconstruction = model(X)
+    reconstruction.square().mean().backward()
+
+    parameters = {
+        "encoder_temporal": (model.encoder_blocks[0].temporal_attention.in_proj_weight),
+        "encoder_spatial": (model.encoder_blocks[0].spatial_attention.in_proj_weight),
+        "decoder_temporal": (model.decoder_blocks[0].temporal_attention.in_proj_weight),
+        "decoder_spatial": (model.decoder_blocks[0].spatial_attention.in_proj_weight),
+        "latent": model.latent.weight,
+        "time_query": model.decoder_time_query,
+        "space_query": model.decoder_space_query,
+    }
+    for name, parameter in parameters.items():
+        assert parameter.grad is not None, name
+        assert torch.isfinite(parameter.grad).all(), name
+        assert torch.count_nonzero(parameter.grad) > 0, name
+
+
+def _set_identity_attention(attention):
+    dimension = attention.embed_dim
+    identity = torch.eye(dimension)
+    with torch.no_grad():
+        attention.in_proj_weight.zero_()
+        attention.in_proj_weight[:dimension].copy_(identity)
+        attention.in_proj_weight[dimension : 2 * dimension].copy_(identity)
+        attention.in_proj_weight[2 * dimension :].copy_(identity)
+        attention.in_proj_bias.zero_()
+        attention.out_proj.weight.copy_(identity)
+        attention.out_proj.bias.zero_()
+
+
+def _zero_module(module):
+    with torch.no_grad():
+        for parameter in module.parameters():
+            parameter.zero_()
+
+
+def test_factorized_block_temporal_attention_mixes_timesteps():
+    block = _FactorizedSpatiotemporalBlock(d_model=4, n_heads=1)
+    _set_identity_attention(block.temporal_attention)
+    _zero_module(block.spatial_attention)
+    _zero_module(block.feed_forward)
+
+    values = torch.zeros(1, 3, 2, 4)
+    values[0, 0, 0] = torch.tensor([1.0, -1.0, 0.5, -0.5])
+    enabled = block(values.clone())
+    with torch.no_grad():
+        block.temporal_attention.out_proj.weight.zero_()
+    disabled = block(values.clone())
+
+    cross_time_change = torch.abs(enabled[0, 1, 0] - disabled[0, 1, 0])
+    assert torch.max(cross_time_change) > 1e-4
+
+
+def test_factorized_block_spatial_attention_mixes_tokens():
+    block = _FactorizedSpatiotemporalBlock(d_model=4, n_heads=1)
+    _zero_module(block.temporal_attention)
+    _set_identity_attention(block.spatial_attention)
+    _zero_module(block.feed_forward)
+
+    values = torch.zeros(1, 2, 3, 4)
+    values[0, 0, 0] = torch.tensor([1.0, -1.0, 0.5, -0.5])
+    enabled = block(values.clone())
+    with torch.no_grad():
+        block.spatial_attention.out_proj.weight.zero_()
+    disabled = block(values.clone())
+
+    cross_space_change = torch.abs(enabled[0, 0, 1] - disabled[0, 0, 1])
+    assert torch.max(cross_space_change) > 1e-4
+
+
+def test_spatial_token_checkpoint_round_trip(tmp_path):
+    X = np.random.randn(12, 3, 1, 8, 8).astype("float32")
+    model = _fit_model(X)
+    checkpoint = tmp_path / "spatial_token.pt"
+
+    original = model.predict(X, verbose=0)
+    model.save_pytorch_model(checkpoint)
+    restored = SpatialTokenConvLSTMTransformerAutoencoder.from_pytorch_model(
+        checkpoint,
+        device="cpu",
+    )
+
+    assert restored.spatial_pool_size == (2, 2)
+    assert np.allclose(original, restored.predict(X, verbose=0))
+
+
+def test_spatial_token_manual_optimization_reduces_loss():
+    coordinate = torch.linspace(-1.0, 1.0, 6)
+    base = coordinate[:, None] + coordinate[None, :]
+    frames = torch.stack([base, base.square(), torch.sin(base)], dim=0)
+    X = frames[None, :, None].repeat(4, 1, 1, 1, 1)
+
+    model = _model()._build_model(tuple(X.shape))
+    optimizer = torch.optim.Adam(model.parameters(), lr=1e-2)
+
+    with torch.no_grad():
+        initial = torch.mean((model(X) - X) ** 2).item()
+    for _ in range(10):
+        optimizer.zero_grad()
+        loss = torch.mean((model(X) - X) ** 2)
+        loss.backward()
+        optimizer.step()
+    with torch.no_grad():
+        final = torch.mean((model(X) - X) ** 2).item()
+
+    assert final < initial
diff --git a/tests/deeplearning/test_variational_autoencoder.py b/tests/deeplearning/test_variational_autoencoder.py
new file mode 100644
index 0000000..d162308
--- /dev/null
+++ b/tests/deeplearning/test_variational_autoencoder.py
@@ -0,0 +1,502 @@
+"""Tests for the dense variational autoencoder."""
+
+import numpy as np
+import pytest
+import torch
+import torch.nn as nn
+
+from bluemath_tk.deeplearning.autoencoders import VariationalAutoencoder
+
+
+@pytest.fixture(autouse=True)
+def _set_seed():
+    previous_threads = torch.get_num_threads()
+    numpy_state = np.random.get_state()
+    torch_state = torch.random.get_rng_state()
+    cuda_states = torch.cuda.get_rng_state_all() if torch.cuda.is_available() else None
+    try:
+        np.random.seed(401)
+        torch.manual_seed(401)
+        torch.set_num_threads(1)
+        yield
+    finally:
+        torch.set_num_threads(previous_threads)
+        np.random.set_state(numpy_state)
+        torch.random.set_rng_state(torch_state)
+        if cuda_states is not None:
+            torch.cuda.set_rng_state_all(cuda_states)
+
+
+def _fit_model(X, beta=0.1, epochs=2, validation_mc_samples=3):
+    model = VariationalAutoencoder(
+        k=3,
+        hidden_dims=[12, 8],
+        beta=beta,
+        validation_mc_samples=validation_mc_samples,
+        device="cpu",
+    )
+    history = model.fit(
+        X,
+        validation_split=0.25,
+        epochs=epochs,
+        batch_size=4,
+        learning_rate=1e-3,
+        patience=max(epochs, 2),
+        verbose=0,
+    )
+    return model, history
+
+
+def test_vae_reconstructs_original_multidimensional_shape():
+    X = np.random.randn(16, 2, 5).astype("float32")
+    model, history = _fit_model(X)
+
+    prediction = model.predict(X, batch_size=4, verbose=0)
+    latent = model.encode(X, batch_size=4, verbose=0)
+    decoded = model.decode(latent, batch_size=4, verbose=0)
+
+    assert prediction.shape == X.shape
+    assert decoded.shape == X.shape
+    assert latent.shape == (16, 3)
+    assert np.allclose(prediction, decoded, rtol=1e-5, atol=1e-6)
+    assert set(history) == {
+        "train_loss",
+        "train_reconstruction_loss",
+        "train_kl_loss",
+        "val_loss",
+        "val_reconstruction_loss",
+        "val_kl_loss",
+        "val_deterministic_reconstruction_loss",
+    }
+
+
+def test_vae_distribution_and_deterministic_inference():
+    X = np.random.randn(16, 7).astype("float32")
+    model, _ = _fit_model(X)
+
+    mu, log_var = model.encode_distribution(X, batch_size=4)
+    first = model.predict(X, batch_size=4, verbose=0)
+    second = model.predict(X, batch_size=4, verbose=0)
+
+    assert mu.shape == (16, 3)
+    assert log_var.shape == (16, 3)
+    assert np.allclose(model.encode(X, verbose=0), mu)
+    assert np.array_equal(first, second)
+    assert np.isfinite(mu).all()
+    assert np.isfinite(log_var).all()
+
+
+def test_vae_stochastic_operations_are_explicit():
+    X = np.random.randn(16, 7).astype("float32")
+    model, _ = _fit_model(X)
+
+    latent_1 = model.sample_latent(X, batch_size=4)
+    latent_2 = model.sample_latent(X, batch_size=4)
+    prediction_1 = model.predict(X, verbose=0, stochastic=True)
+    prediction_2 = model.predict(X, verbose=0, stochastic=True)
+
+    assert latent_1.shape == (16, 3)
+    assert not np.array_equal(latent_1, latent_2)
+    assert not np.array_equal(prediction_1, prediction_2)
+
+
+def test_vae_kl_has_known_analytic_values():
+    inner = VariationalAutoencoder(
+        k=3,
+        hidden_dims=[4],
+        device="cpu",
+    )._build_model((4, 5))
+
+    zero_mu = torch.zeros(4, 3)
+    zero_log_var = torch.zeros(4, 3)
+    assert inner.kl_divergence(zero_mu, zero_log_var).item() == pytest.approx(0.0)
+
+    unit_mu = torch.ones(4, 3)
+    expected = 0.5 * 3
+    assert inner.kl_divergence(unit_mu, zero_log_var).item() == pytest.approx(expected)
+
+
+def test_vae_kl_is_stable_at_float32_mean_boundary():
+    outer = VariationalAutoencoder(k=1, hidden_dims=[2], device="cpu")
+    inner = outer._build_model((2, 1))
+    maximum = np.finfo(np.float32).max
+    mu = torch.tensor([[maximum]], dtype=torch.float32, requires_grad=True)
+    log_var = torch.zeros_like(mu)
+
+    loss = inner.kl_divergence(mu, log_var)
+    loss.backward()
+
+    assert loss.dtype == torch.float64
+    assert torch.isfinite(loss)
+    assert mu.grad is not None
+    assert torch.isfinite(mu.grad).all()
+    assert torch.count_nonzero(mu.grad) > 0
+
+
+def test_vae_reparameterization_matches_standard_normal_moments():
+    inner = VariationalAutoencoder(
+        k=1,
+        hidden_dims=[4],
+        device="cpu",
+    )._build_model((4, 2))
+    mu = torch.full((20000, 1), 2.0)
+    log_var = torch.log(torch.full((20000, 1), 9.0))
+
+    sample = inner.reparameterize(mu, log_var)
+
+    assert sample.mean().item() == pytest.approx(2.0, abs=0.08)
+    assert sample.std(unbiased=False).item() == pytest.approx(3.0, abs=0.08)
+
+
+def test_vae_full_objective_gradients_reach_all_paths():
+    outer = VariationalAutoencoder(
+        k=2,
+        hidden_dims=[6],
+        beta=0.2,
+        device="cpu",
+    )
+    inner = outer._build_model((8, 4))
+    X = torch.randn(8, 4)
+
+    mu, log_var = inner.encode_distribution_forward(X)
+    reconstruction = inner.decode_forward(inner.reparameterize(mu, log_var))
+    loss = nn.functional.mse_loss(reconstruction, X)
+    loss = loss + outer.beta * inner.kl_divergence(mu, log_var)
+    loss.backward()
+
+    for parameter in (
+        inner.mu_layer.weight,
+        inner.variance_layer.weight,
+        inner.decoder[-1].weight,
+    ):
+        assert parameter.grad is not None
+        assert torch.isfinite(parameter.grad).all()
+        assert torch.count_nonzero(parameter.grad) > 0
+
+
+@pytest.mark.parametrize(
+    "raw_bias",
+    [-1000.0, 1000.0, np.finfo(np.float32).max],
+)
+def test_vae_extreme_variance_logits_keep_corrective_gradients(raw_bias):
+    outer = VariationalAutoencoder(
+        k=2,
+        hidden_dims=[4],
+        beta=1.0,
+        device="cpu",
+    )
+    inner = outer._build_model((6, 3))
+    with torch.no_grad():
+        inner.variance_layer.weight.zero_()
+        inner.variance_layer.bias.fill_(raw_bias)
+
+    X = torch.zeros(6, 3)
+    mu, log_var = inner.encode_distribution_forward(X)
+    loss = inner.kl_divergence(mu, log_var)
+    loss.backward()
+
+    gradient = inner.variance_layer.bias.grad
+    assert torch.isfinite(log_var).all()
+    assert torch.isfinite(loss)
+    assert gradient is not None
+    assert torch.isfinite(gradient).all()
+    assert torch.count_nonzero(gradient) > 0
+
+
+def test_vae_validation_reports_stochastic_and_deterministic_metrics():
+    outer = VariationalAutoencoder(
+        k=1,
+        hidden_dims=[2],
+        beta=0.0,
+        validation_mc_samples=64,
+        device="cpu",
+    )
+
+    class KnownPosterior(nn.Module):
+        def __init__(self):
+            super().__init__()
+            self.anchor = nn.Parameter(torch.zeros(()))
+
+        def encode_distribution_forward(self, x):
+            mu = torch.zeros(len(x), 1) + self.anchor * 0
+            log_var = torch.zeros_like(mu)
+            return mu, log_var
+
+        @staticmethod
+        def reparameterize(mu, log_var):
+            return mu + torch.exp(0.5 * log_var) * torch.randn_like(mu)
+
+        @staticmethod
+        def kl_divergence(mu, log_var):
+            return (
+                0.5
+                * torch.sum(
+                    mu.pow(2) + log_var.exp() - 1.0 - log_var,
+                    dim=1,
+                ).mean()
+            )
+
+        @staticmethod
+        def decode_forward(z):
+            return torch.relu(z)
+
+    outer.model = KnownPosterior()
+    X = torch.zeros(32, 1)
+    totals = outer._run_vae_epoch(
+        X,
+        X,
+        batch_size=8,
+        criterion=nn.MSELoss(),
+        optimizer=None,
+        stochastic_samples=64,
+        report_deterministic=True,
+    )
+
+    assert totals["deterministic_reconstruction_loss"] == pytest.approx(0.0)
+    assert totals["reconstruction_loss"] > 0.2
+    assert totals["loss"] == pytest.approx(totals["reconstruction_loss"])
+
+
+def test_vae_early_stopping_uses_stochastic_validation_objective(monkeypatch):
+    X = np.zeros((8, 3), dtype="float32")
+    model = VariationalAutoencoder(
+        k=2,
+        hidden_dims=[4],
+        beta=0.1,
+        validation_mc_samples=5,
+        device="cpu",
+    )
+    validation_losses = iter([3.0, 1.0, 2.0])
+    training_epoch = {"value": 0}
+
+    def fake_epoch(
+        X_tensor,
+        y_tensor,
+        batch_size,
+        criterion,
+        optimizer,
+        stochastic_samples,
+        report_deterministic,
+    ):
+        if optimizer is not None:
+            training_epoch["value"] += 1
+            with torch.no_grad():
+                model.model.mu_layer.bias.fill_(training_epoch["value"])
+            return {
+                "loss": 1.0,
+                "reconstruction_loss": 1.0,
+                "kl_loss": 0.0,
+                "deterministic_reconstruction_loss": 0.0,
+            }
+
+        assert stochastic_samples == model.validation_mc_samples
+        assert report_deterministic
+        validation_loss = next(validation_losses)
+        return {
+            "loss": validation_loss,
+            "reconstruction_loss": validation_loss,
+            "kl_loss": 0.0,
+            "deterministic_reconstruction_loss": 100.0 - validation_loss,
+        }
+
+    monkeypatch.setattr(model, "_run_vae_epoch", fake_epoch)
+    model.fit(
+        X,
+        validation_split=0.25,
+        epochs=3,
+        batch_size=4,
+        patience=3,
+        verbose=0,
+    )
+
+    expected = torch.full_like(model.model.mu_layer.bias, 2.0)
+    assert torch.allclose(model.model.mu_layer.bias, expected)
+
+
+def test_vae_validation_sampling_does_not_advance_training_rng(monkeypatch):
+    X = np.zeros((8, 3), dtype="float32")
+    model = VariationalAutoencoder(
+        k=2,
+        hidden_dims=[4],
+        validation_mc_samples=7,
+        device="cpu",
+    )
+    model.model = model._build_model(X.shape).to(model.device)
+    optimizer = torch.optim.SGD(model.model.parameters(), lr=0.0)
+    training_draws = []
+
+    def fake_epoch(
+        X_tensor,
+        y_tensor,
+        batch_size,
+        criterion,
+        optimizer,
+        stochastic_samples,
+        report_deterministic,
+    ):
+        if optimizer is not None:
+            training_draws.append(torch.rand(()).item())
+        else:
+            torch.rand(128)
+        return {
+            "loss": 1.0,
+            "reconstruction_loss": 1.0,
+            "kl_loss": 0.0,
+            "deterministic_reconstruction_loss": 1.0,
+        }
+
+    monkeypatch.setattr(model, "_run_vae_epoch", fake_epoch)
+    torch.manual_seed(811)
+    model.fit(
+        X,
+        validation_split=0.25,
+        epochs=3,
+        batch_size=4,
+        optimizer=optimizer,
+        patience=3,
+        verbose=0,
+    )
+
+    torch.manual_seed(811)
+    expected_draws = [torch.rand(()).item() for _ in range(3)]
+    assert training_draws == pytest.approx(expected_draws)
+
+
+def test_vae_beta_zero_removes_kl_from_total_loss():
+    X = np.random.randn(16, 6).astype("float32")
+    _, history = _fit_model(X, beta=0.0, epochs=1)
+
+    assert history["train_loss"] == pytest.approx(history["train_reconstruction_loss"])
+    assert history["val_loss"] == pytest.approx(history["val_reconstruction_loss"])
+
+
+def test_vae_custom_inference_rejects_float32_overflow():
+    X = np.zeros((16, 6), dtype="float32")
+    model, _ = _fit_model(X)
+    huge = np.full((2, 6), 1e100, dtype="float64")
+
+    for method in (
+        model.encode_distribution,
+        model.sample_latent,
+    ):
+        with pytest.raises(ValueError, match="converted to float32"):
+            method(huge)
+    with pytest.raises(ValueError, match="converted to float32"):
+        model.predict(huge, verbose=0, stochastic=True)
+
+
+def test_vae_custom_inference_rejects_nonfinite_model_outputs(
+    monkeypatch,
+):
+    X = np.zeros((16, 6), dtype="float32")
+    model, _ = _fit_model(X)
+
+    def nonfinite_distribution(batch):
+        mu = torch.full(
+            (len(batch), model.k),
+            float("inf"),
+            dtype=batch.dtype,
+            device=batch.device,
+        )
+        return mu, torch.zeros_like(mu)
+
+    monkeypatch.setattr(
+        model.model,
+        "encode_distribution_forward",
+        nonfinite_distribution,
+    )
+    for method in (
+        model.encode_distribution,
+        model.sample_latent,
+    ):
+        with pytest.raises(FloatingPointError, match="not finite"):
+            method(X[:2])
+    with pytest.raises(FloatingPointError, match="not finite"):
+        model.predict(X[:2], verbose=0, stochastic=True)
+
+    sample_shape = tuple(X.shape[1:])
+
+    def nonfinite_decode(z):
+        return torch.full(
+            (len(z), *sample_shape),
+            float("inf"),
+            dtype=z.dtype,
+            device=z.device,
+        )
+
+    monkeypatch.setattr(model.model, "decode_forward", nonfinite_decode)
+    with pytest.raises(FloatingPointError, match="not finite"):
+        model.sample(2)
+
+
+def test_vae_sampling_and_single_vector_decoding():
+    X = np.random.randn(16, 6).astype("float32")
+    model, _ = _fit_model(X)
+    latent = model.encode(X, verbose=0)
+
+    decoded = model.decode(latent[0], verbose=0)
+    generated = model.sample(5, batch_size=2)
+
+    assert decoded.shape == (1, 6)
+    assert generated.shape == (5, 6)
+    assert np.isfinite(generated).all()
+
+
+def test_vae_checkpoint_round_trip(tmp_path):
+    X = np.random.randn(16, 6).astype("float32")
+    model, _ = _fit_model(X)
+    checkpoint = tmp_path / "vae.pt"
+
+    original = model.predict(X, verbose=0)
+    model.save_pytorch_model(checkpoint)
+    restored = VariationalAutoencoder.from_pytorch_model(
+        checkpoint,
+        device="cpu",
+    )
+
+    assert restored.beta == model.beta
+    assert restored.hidden_dims == model.hidden_dims
+    assert restored.validation_mc_samples == model.validation_mc_samples
+    assert np.allclose(original, restored.predict(X, verbose=0))
+
+
+@pytest.mark.parametrize(
+    ("kwargs", "message"),
+    [
+        ({"k": 0}, "k"),
+        ({"hidden_dims": []}, "hidden_dims"),
+        ({"hidden_dims": [8, 0]}, "hidden dimension"),
+        ({"beta": -1.0}, "beta"),
+        ({"beta": float("nan")}, "beta"),
+        ({"beta": float("inf")}, "beta"),
+        ({"validation_mc_samples": 0}, "validation_mc_samples"),
+    ],
+)
+def test_vae_rejects_invalid_constructor_arguments(kwargs, message):
+    with pytest.raises(ValueError, match=message):
+        VariationalAutoencoder(**kwargs)
+
+
+def test_vae_rejects_invalid_sampling_requests():
+    model = VariationalAutoencoder(k=2, hidden_dims=[4], device="cpu")
+
+    with pytest.raises(ValueError, match="fitted"):
+        model.sample(1)
+    with pytest.raises(TypeError, match="integer"):
+        model.sample(1.5)
+    with pytest.raises(TypeError, match="batch_size"):
+        model.sample(1, batch_size=1.5)
+
+
+def test_vae_training_reduces_structured_reconstruction_loss():
+    coordinate = np.linspace(-1.0, 1.0, 32, dtype="float32")[:, None]
+    X = np.concatenate(
+        [coordinate, coordinate**2, np.sin(np.pi * coordinate)],
+        axis=1,
+    )
+    _, history = _fit_model(X, beta=0.0, epochs=12)
+
+    assert (
+        history["train_reconstruction_loss"][-1]
+        < (history["train_reconstruction_loss"][0])
+    )