Add variational and spatial token autoencoders

[sergioald]

Summary

This PR extends the BlueMath_tk autoencoder framework with two advanced model
families:

• VariationalAutoencoder
• SpatialTokenConvLSTMTransformerAutoencoder

It also strengthens the shared autoencoder infrastructure with stricter input
validation, transactional checkpoint loading, numerically stable scientific
reconstruction metrics, consistent latent-shape handling, and expanded
regression coverage.

This branch is based on the full-sequence reconstruction behaviour introduced
in PR #154.

New models

VariationalAutoencoder

The new variational autoencoder provides:

• a configurable beta-VAE objective;
• a diagonal-Gaussian latent distribution;
• reparameterised latent sampling;
• numerically stable variance handling;
• stochastic Monte Carlo validation;
• separate deterministic reconstruction diagnostics;
• multidimensional input reconstruction;
• fit, predict, encode, decode, evaluate, and checkpoint support.

The implemented objective is:

reconstruction loss + beta * KL(q(z|x) || N(0, I))

Validation uses the same stochastic beta-VAE objective used during training.
Posterior-mean reconstruction is reported separately through
val_deterministic_reconstruction_loss.

SpatialTokenConvLSTMTransformerAutoencoder

The new spatial-token model operates on full spatiotemporal sequences:

(B, T, C, H, W) -> (B, k) -> (B, T, C, H, W)

Its architecture combines:

• ConvLSTM-based spatiotemporal feature extraction;
• temporal attention applied independently to each spatial token;
• spatial attention applied independently at each time step;
• one strict latent bottleneck of size k;
• full-sequence decoding;
• support for odd and non-divisible spatial dimensions.

The decoder receives only the latent vector and learned model parameters. There
are no skip connections, cached inputs, or encoder feature maps that bypass the
latent bottleneck.

Shared autoencoder hardening

The PR improves validation and failure handling across the existing and new
autoencoder families.

Input and target validation

• Reconstruction targets must match the input shape exactly.
• Broadcastable but incompatible targets are rejected.
• Prediction and encoding inputs must match the fitted per-sample shape.
• Decoder inputs must be either:
  • one latent vector with shape (k,); or
  • a latent batch with shape (B, k).
• Empty, non-numeric, complex, NaN, infinite, and float32-overflowing arrays are
  rejected.
• Invalid learning rates and training arguments are rejected before model or
  build-state mutation.

Non-finite execution protection

Training, validation, inference, and checkpoint operations check for non-finite:

• losses;
• model outputs;
• gradients;
• parameters;
• floating and complex buffers.

A failed fit cannot leave a model marked as successfully fitted with invalid
state.

Transactional checkpoint loading

Checkpoint loading is staged before modifying a live model.

A failed checkpoint load preserves:

• the existing model object;
• parameters and buffers;
• fitted state;
• build input shape;
• constructor configuration;
• checkpoint-derived metadata.

Checkpoints containing values that are finite in their stored dtype but overflow
when converted to the destination dtype are rejected before the live model is
modified.

Scientific reconstruction metrics

The reconstruction metric API now provides robust NumPy and PyTorch
implementations of:

• MSE;
• MAE;
• RMSE.

Supported reductions are:

• none;
• sample;
• mean;
• sum.

The implementation includes:

• exact shape validation;
• finite-input validation;
• validation of eps;
• stable float32 and float64 calculations;
• exact-zero RMSE with finite zero gradients;
• protection against intermediate underflow and overflow;
• stable averaging of multiple extreme sample errors;
• detection of genuinely non-representable final results;
• preservation of PyTorch dtype, device, and autograd behaviour.

Very small non-zero residuals remain distinguishable from exact reconstruction,
while very large but representable residuals do not fail merely because their
direct squares or intermediate sums overflow.

Test-state isolation

The advanced autoencoder tests now restore:

• NumPy random-number state;
• PyTorch CPU random-number state;
• CUDA random-number state when CUDA is available;
• the previous PyTorch thread count.

Restoration is checked after both successful and deliberately failing inner test
runs.

Public API

The intended public autoencoder exports are explicitly declared:

• StandardAutoencoder
• OrthogonalAutoencoder
• LSTMAutoencoder
• CNNAutoencoder
• VisionTransformerAutoencoder
• ConvLSTMAutoencoder
• HybridConvLSTMTransformerAutoencoder
• VariationalAutoencoder
• SpatialTokenConvLSTMTransformerAutoencoder

The common workflow is:

model.fit(X)

latent = model.encode(X)
prediction = model.predict(X)
decoded = model.decode(latent)
metrics = model.evaluate(X)

For the full-sequence spatiotemporal models:

assert latent.shape == (len(X), model.k)
assert prediction.shape == X.shape
assert decoded.shape == X.shape

Scope and limitations

The latent vector provides dimensional compression, but this PR does not
implement entropy coding, bitrate control, latent quantisation, or a deployable
storage codec.

Other current limitations include:

• inputs are converted to float32 by the model API;
• complete training and validation splits are transferred to the selected
  device;
• full-resolution ConvLSTM activations can require substantial memory;
• beta=0 does not produce a prior-matched generative model;
• CUDA-specific behaviour could not be validated locally because the test
  environment did not include a CUDA-capable device.

Chronological experiment splitting, common benchmarking, latent quantisation,
storage accounting, and downstream scientific evaluation are intended for
separate follow-up work.

Tests

The added and expanded tests cover:

• VAE objective and KL-divergence correctness;
• reparameterisation;
• stable variance gradients;
• stochastic and deterministic VAE validation;
• full-sequence reconstruction;
• temporal and spatial attention behaviour;
• strict latent bottlenecks;
• malformed latent inputs across all model families;
• exact target and output shape validation;
• short final-batch weighting;
• non-finite losses, gradients, parameters, buffers, and outputs;
• transactional checkpoint failure handling;
• constructor and learning-rate validation;
• extreme float32 and float64 reconstruction metrics;
• exact-zero RMSE gradients;
• public exports;
• test RNG and thread-state isolation.

Local validation

Post-rebase validation against the current develop branch:

compileall: passed

focused Ruff checks:
- _base_model.py: passed
- autoencoders.py: passed
- layers.py: passed
- metrics.py: passed
- variational_autoencoders.py: passed
- spatiotemporal_autoencoders.py: passed
- advanced test files: passed

Ruff formatting:
- 11 files already formatted

deep-learning suite:
- 357 passed
- 88 warnings