pyautoencoder 1.0.4__tar.gz → 1.0.6__tar.gz

pyautoencoder-1.0.6/PKG-INFO

@@ -0,0 +1,162 @@
+Metadata-Version: 2.4
+Name: pyautoencoder
+Version: 1.0.6
+Summary: A Python package offering implementations of state-of-the-art autoencoder architectures in PyTorch.
+Home-page: https://github.com/andrea-pollastro/pyautoencoder
+Author: Andrea Pollastro
+License: MIT
+Keywords: autoencoder,pytorch,deep learning,machine learning,representation learning,dimensionality reduction,generative models
+Classifier: Operating System :: OS Independent
+Classifier: Intended Audience :: Science/Research
+Classifier: Intended Audience :: Developers
+Classifier: License :: OSI Approved :: MIT License
+Classifier: Programming Language :: Python
+Classifier: Topic :: Scientific/Engineering :: Artificial Intelligence
+Classifier: Topic :: Software Development :: Libraries :: Python Modules
+Requires-Python: >=3.7
+Description-Content-Type: text/markdown
+License-File: LICENSE
+Requires-Dist: torch>=2.0.0
+Dynamic: author
+Dynamic: classifier
+Dynamic: description
+Dynamic: description-content-type
+Dynamic: home-page
+Dynamic: keywords
+Dynamic: license
+Dynamic: license-file
+Dynamic: requires-dist
+Dynamic: requires-python
+Dynamic: summary
+
+# PyAutoencoder
+
+A clean, modular PyTorch library for building and training autoencoders.
+
+<!-- <p align="center">
+  <img src="assets/logo_nobackground.png" alt="pyautoencoder logo" width="220"/>
+</p> -->
+
+![logo](https://raw.githubusercontent.com/andrea-pollastro/pyautoencoder/main/assets/logo_nobackground.png)
+
+<p align="center">
+  <a href="https://pypi.org/project/pyautoencoder/"><img alt="PyPI" src="https://img.shields.io/pypi/v/pyautoencoder.svg"></a>
+  <a href="https://github.com/andrea-pollastro/pyautoencoder/blob/main/LICENSE"><img alt="License: MIT" src="https://img.shields.io/badge/License-MIT-blue.svg"></a>
+  <a href="https://github.com/andrea-pollastro/pyautoencoder/stargazers"><img alt="Stars" src="https://img.shields.io/github/stars/andrea-pollastro/pyautoencoder?style=social"></a>
+</p>
+
+---
+
+## Highlights
+
+PyAutoencoder is designed to offer **simple and easy access to autoencoder frameworks**. Here's what it offers:
+
+- **Minimal, composable API**  
+You don't have to inherit from complicated base classes or learn a new training loop. Simply provide your own PyTorch nn.Module encoder and decoder, and plug them into the ready‑to‑use autoencoder wrappers. This makes it easy to experiment with different architectures (e.g. MLPs, CNNs) while reusing the same training pipeline.
+
+- **Ready‑to‑use autoencoders**
+  The library ships with working implementations of autoencoders, each paired with their respective loss functions. You can start training in a few lines, without re‑implementing reconstruction likelihoods, KL divergence, or other boilerplate.
+
+- **PyTorch compatibility**  
+  The library is fully compatible with the PyTorch ecosystem, so models integrate naturally with modules, tensors, optimizers, and schedulers.
+
+- **Lightweight, research‑oriented**  
+  The library is intentionally minimal: no training loop frameworks, no heavy abstractions. This makes it well suited for research prototypes where you want control and transparency.
+
+> **Status**: The project is in an early but usable stage. Contributions, issues, and feedback are highly encouraged!
+
+**Currently implemented**:
+- Autoencoder (AE)
+- Variational Autoencoder (VAE)
+
+---
+
+## Installation
+
+```bash
+pip install pyautoencoder
+```
+
+Or install from source for development:
+
+```bash
+git clone https://github.com/andrea-pollastro/pyautoencoder.git
+cd pyautoencoder
+pip install -e .
+```
+
+## Quick start
+
+```python
+import torch
+import torch.nn as nn
+from pyautoencoder import VAE, VAELoss
+
+# Define encoder/decoder
+encoder = nn.Sequential(
+    nn.Linear(784, 512), 
+    nn.ReLU(),
+    nn.Linear(512, 256)
+)
+
+decoder = nn.Sequential(
+    nn.Linear(256, 512), 
+    nn.ReLU(),
+    nn.Linear(512, 784)
+)
+
+# Model
+vae = VAE(encoder=encoder, decoder=decoder, latent_dim=32)
+
+# Loss
+criterion = VAELoss(beta=1.0, likelihood="gaussian")
+optimizer = torch.optim.Adam(vae.parameters())
+for x in dataloader:
+    optimizer.zero_grad()
+    out = vae(x)
+    losses = criterion(out, x)
+    losses.total.backward() # negative ELBO
+    optimizer.step()
+
+    # optional: log components
+    log_likelihood = losses.components["log_likelihood"]
+    kl_divergence = losses.components["kl_divergence"]
+```
+
+## Built‑in models
+
+- **`AE`** — standard Autoencoder
+  ```python
+  from pyautoencoder import AE, AutoencoderLoss
+  ae = AE(encoder=encoder, decoder=decoder)
+  criterion = AutoencoderLoss(likelihood="gaussian") # or bernoulli
+  ```
+
+- **`VAE`** — Variational Autoencoder
+  ```python
+  from pyautoencoder import VAE, VAELoss
+  vae = VAE(encoder=encoder, decoder=decoder, latent_dim=32)
+  criterion = VAELoss(beta=1.0, likelihood="gaussian") # or bernoulli
+  ```
+
+## Examples
+
+See the [`examples/`](examples/) folder for runnable scripts showing example of usage.
+
+## License
+
+This project is released under the **MIT License**. See [LICENSE](LICENSE).
+
+## Citation
+
+If you use this package in academic work, please cite:
+
+```bibtex
+@misc{pollastro2025pyautoencoder,
+  author       = {Andrea Pollastro},
+  title        = {pyautoencoder},
+  year         = {2025},
+  howpublished = {GitHub repository},
+  url          = {https://github.com/andrea-pollastro/pyautoencoder}
+}
+```

pyautoencoder-1.0.6/README.md

@@ -0,0 +1,131 @@
+# PyAutoencoder
+
+A clean, modular PyTorch library for building and training autoencoders.
+
+<!-- <p align="center">
+  <img src="assets/logo_nobackground.png" alt="pyautoencoder logo" width="220"/>
+</p> -->
+
+![logo](https://raw.githubusercontent.com/andrea-pollastro/pyautoencoder/main/assets/logo_nobackground.png)
+
+<p align="center">
+  <a href="https://pypi.org/project/pyautoencoder/"><img alt="PyPI" src="https://img.shields.io/pypi/v/pyautoencoder.svg"></a>
+  <a href="https://github.com/andrea-pollastro/pyautoencoder/blob/main/LICENSE"><img alt="License: MIT" src="https://img.shields.io/badge/License-MIT-blue.svg"></a>
+  <a href="https://github.com/andrea-pollastro/pyautoencoder/stargazers"><img alt="Stars" src="https://img.shields.io/github/stars/andrea-pollastro/pyautoencoder?style=social"></a>
+</p>
+
+---
+
+## Highlights
+
+PyAutoencoder is designed to offer **simple and easy access to autoencoder frameworks**. Here's what it offers:
+
+- **Minimal, composable API**  
+You don't have to inherit from complicated base classes or learn a new training loop. Simply provide your own PyTorch nn.Module encoder and decoder, and plug them into the ready‑to‑use autoencoder wrappers. This makes it easy to experiment with different architectures (e.g. MLPs, CNNs) while reusing the same training pipeline.
+
+- **Ready‑to‑use autoencoders**
+  The library ships with working implementations of autoencoders, each paired with their respective loss functions. You can start training in a few lines, without re‑implementing reconstruction likelihoods, KL divergence, or other boilerplate.
+
+- **PyTorch compatibility**  
+  The library is fully compatible with the PyTorch ecosystem, so models integrate naturally with modules, tensors, optimizers, and schedulers.
+
+- **Lightweight, research‑oriented**  
+  The library is intentionally minimal: no training loop frameworks, no heavy abstractions. This makes it well suited for research prototypes where you want control and transparency.
+
+> **Status**: The project is in an early but usable stage. Contributions, issues, and feedback are highly encouraged!
+
+**Currently implemented**:
+- Autoencoder (AE)
+- Variational Autoencoder (VAE)
+
+---
+
+## Installation
+
+```bash
+pip install pyautoencoder
+```
+
+Or install from source for development:
+
+```bash
+git clone https://github.com/andrea-pollastro/pyautoencoder.git
+cd pyautoencoder
+pip install -e .
+```
+
+## Quick start
+
+```python
+import torch
+import torch.nn as nn
+from pyautoencoder import VAE, VAELoss
+
+# Define encoder/decoder
+encoder = nn.Sequential(
+    nn.Linear(784, 512), 
+    nn.ReLU(),
+    nn.Linear(512, 256)
+)
+
+decoder = nn.Sequential(
+    nn.Linear(256, 512), 
+    nn.ReLU(),
+    nn.Linear(512, 784)
+)
+
+# Model
+vae = VAE(encoder=encoder, decoder=decoder, latent_dim=32)
+
+# Loss
+criterion = VAELoss(beta=1.0, likelihood="gaussian")
+optimizer = torch.optim.Adam(vae.parameters())
+for x in dataloader:
+    optimizer.zero_grad()
+    out = vae(x)
+    losses = criterion(out, x)
+    losses.total.backward() # negative ELBO
+    optimizer.step()
+
+    # optional: log components
+    log_likelihood = losses.components["log_likelihood"]
+    kl_divergence = losses.components["kl_divergence"]
+```
+
+## Built‑in models
+
+- **`AE`** — standard Autoencoder
+  ```python
+  from pyautoencoder import AE, AutoencoderLoss
+  ae = AE(encoder=encoder, decoder=decoder)
+  criterion = AutoencoderLoss(likelihood="gaussian") # or bernoulli
+  ```
+
+- **`VAE`** — Variational Autoencoder
+  ```python
+  from pyautoencoder import VAE, VAELoss
+  vae = VAE(encoder=encoder, decoder=decoder, latent_dim=32)
+  criterion = VAELoss(beta=1.0, likelihood="gaussian") # or bernoulli
+  ```
+
+## Examples
+
+See the [`examples/`](examples/) folder for runnable scripts showing example of usage.
+
+## License
+
+This project is released under the **MIT License**. See [LICENSE](LICENSE).
+
+## Citation
+
+If you use this package in academic work, please cite:
+
+```bibtex
+@misc{pollastro2025pyautoencoder,
+  author       = {Andrea Pollastro},
+  title        = {pyautoencoder},
+  year         = {2025},
+  howpublished = {GitHub repository},
+  url          = {https://github.com/andrea-pollastro/pyautoencoder}
+}
+```

pyautoencoder-1.0.6/pyautoencoder/__init__.py

@@ -0,0 +1,15 @@
+"""PyAutoencoder: A clean, modular PyTorch library for autoencoder models."""
+
+from .models.autoencoder import AE
+from .models.variational.vae import VAE
+from .loss.wrapper import AELoss, VAELoss, LossComponents
+
+__version__ = "0.1.0"
+
+__all__ = [
+    'AE',
+    'VAE',
+    'AELoss',
+    'VAELoss',
+    'LossComponents'
+]

pyautoencoder-1.0.6/pyautoencoder/loss/__init__.py

@@ -0,0 +1,12 @@
+"""Loss functions and wrappers for autoencoders."""
+from .wrapper import AELoss, VAELoss, LossComponents
+from .base import log_likelihood
+from .vae import kl_divergence_gaussian
+
+__all__ = [
+    'AELoss',
+    'VAELoss',
+    'LossComponents',
+    'log_likelihood',
+    'kl_divergence_gaussian'
+]

pyautoencoder-1.0.6/pyautoencoder/loss/base.py

@@ -0,0 +1,67 @@
+"""Base loss functions for autoencoders."""
+import math
+import torch
+import torch.nn.functional as F
+from typing import Union
+from enum import Enum
+
+class LikelihoodType(Enum):
+    """Supported likelihood types for the decoder distribution p(x|z)."""
+    GAUSSIAN = 'gaussian'
+    BERNOULLI = 'bernoulli'
+
+# Cache for log(2π) constants per (device, dtype)
+_LOG2PI_CACHE = {}
+
+def _get_log2pi(x: torch.Tensor) -> torch.Tensor:
+    """Return log(2π) cached for the given device/dtype."""
+    key = (x.device, x.dtype)
+    if key not in _LOG2PI_CACHE:
+        _LOG2PI_CACHE[key] = torch.tensor(2.0 * math.pi, device=x.device, dtype=x.dtype).log()
+    return _LOG2PI_CACHE[key]
+
+def log_likelihood(x: torch.Tensor, 
+                   x_hat: torch.Tensor, 
+                   likelihood: Union[str, LikelihoodType] = LikelihoodType.GAUSSIAN) -> torch.Tensor:
+    """
+    Computes elementwise log-likelihood log p(x|x_hat) under different likelihood assumptions.
+    
+    For continuous data:
+        Gaussian (σ² = 1):
+            log p(x|x_hat) = -0.5 * [ (x - x_hat)^2 + log(2π) ]
+        Each dimension contributes independently. To obtain per-sample log-likelihoods,
+        sum over feature dimensions.
+    
+    For discrete data:
+        Bernoulli:
+            log p(x|x_hat) = x * log σ(x_hat) + (1 - x) * log(1 - σ(x_hat)),
+        where σ is the sigmoid function and x_hat are logits.
+    
+    Args:
+        x (torch.Tensor): Ground truth tensor.
+        x_hat (torch.Tensor): Reconstructed tensor. For Bernoulli, values are logits.
+        likelihood (Union[str, LikelihoodType]): Choice of likelihood model. Defaults to Gaussian.
+    
+    Returns:
+        torch.Tensor: Elementwise log-likelihood with the same shape as `x`.
+                      For multi-dimensional inputs, reduce across feature dimensions
+                      to obtain per-sample log-likelihoods.
+    
+    Notes:
+        - Bernoulli case uses a numerically stable BCE implementation in log-space.
+        - Gaussian case assumes fixed unit variance (σ²=1) and includes the normalization constant.
+        - log(2π) is cached per (device, dtype) for efficiency.
+    """
+    if isinstance(likelihood, str):
+        likelihood = LikelihoodType(likelihood.lower())
+    
+    if likelihood == LikelihoodType.BERNOULLI:
+        return -F.binary_cross_entropy_with_logits(x_hat, x, reduction='none')
+    
+    elif likelihood == LikelihoodType.GAUSSIAN:
+        squared_error = (x_hat - x) ** 2
+        log_2pi = _get_log2pi(x)
+        return -0.5 * (squared_error + log_2pi)
+    
+    else:
+        raise ValueError(f"Unsupported likelihood: {likelihood}")

pyautoencoder-1.0.6/pyautoencoder/loss/vae.py

@@ -0,0 +1,88 @@
+"""Loss functions for variational autoencoders with rigorous mathematical implementations."""
+import torch
+from typing import Union, NamedTuple
+
+from .base import log_likelihood, LikelihoodType
+
+class ELBOComponents(NamedTuple):
+    """Components of the ELBO computation."""
+    elbo: torch.Tensor                 # scalar: batch-mean ELBO (with grad)
+    log_likelihood: torch.Tensor       # scalar: batch-mean E_q[log p(x|z)]
+    beta_kl_divergence: torch.Tensor   # scalar: batch-mean β * KL(q||p)
+
+def kl_divergence_gaussian(mu: torch.Tensor, log_var: torch.Tensor) -> torch.Tensor:
+    """
+    Computes the KL divergence KL(q(z|x) || p(z)) between the approximate
+    posterior q(z|x) = N(μ, σ²) and the standard normal prior p(z) = N(0, I).
+
+    Args:
+        mu (torch.Tensor): Mean of q(z|x), shape [B, D_z].
+        log_var (torch.Tensor): Log-variance of q(z|x), shape [B, D_z].
+
+    Returns:
+        torch.Tensor: KL divergence per sample, shape [B].
+                      Reduction over latent dimensions is performed inside,
+                      but not over the batch.
+    """
+    return -0.5 * torch.sum(1 + log_var - mu.pow(2) - log_var.exp(), dim=-1)
+
+def compute_ELBO(
+    x: torch.Tensor,
+    x_hat: torch.Tensor,
+    mu: torch.Tensor,
+    log_var: torch.Tensor,
+    likelihood: Union[str, LikelihoodType] = LikelihoodType.GAUSSIAN,
+    beta: float = 1.0,
+) -> ELBOComponents:
+    """
+    Computes the Evidence Lower Bound (ELBO) for a Variational Autoencoder
+    using the β-VAE formulation.
+
+    Args:
+        x (torch.Tensor): Ground truth inputs, shape [B, ...].
+        x_hat (torch.Tensor): Reconstructed samples, shape [B, ...] or [B, S, ...],
+                              where S is the number of Monte Carlo samples from q(z|x).
+        mu (torch.Tensor): Mean of q(z|x), shape [B, D_z].
+        log_var (torch.Tensor): Log-variance of q(z|x), shape [B, D_z].
+        likelihood (Union[str, LikelihoodType]): Likelihood model for p(x|z).
+        beta (float): Weighting factor for the KL term (β-VAE).
+
+    Returns:
+        ELBOComponents: NamedTuple containing:
+            - elbo (torch.Tensor): Scalar, mean ELBO over the batch.
+            - log_likelihood (torch.Tensor): Scalar, mean reconstruction term
+              E_q[log p(x|z)] over the batch.
+            - beta_kl_divergence (torch.Tensor): Scalar, β * mean KL divergence over the batch.
+
+    Notes:
+        - If x_hat has no sample dimension, it is assumed to contain a single sample (S=1).
+        - log p(x|z) is computed using the `log_likelihood` function, which already
+          includes the Gaussian normalization constant for σ²=1 or the stable BCE for Bernoulli.
+        - All outputs are averaged over the batch for reporting and optimization.
+    """
+    # Ensure a sample dimension S exists -> [B, S, ...]
+    if x_hat.dim() == x.dim():
+        x_hat = x_hat.unsqueeze(1)  # S = 1
+    B, S = x_hat.size(0), x_hat.size(1)
+
+    # log p(x|z): elementwise -> sum over features => [B, S]
+    log_px_z = log_likelihood(x.unsqueeze(1), x_hat, likelihood=likelihood)
+    log_px_z = log_px_z.reshape(B, S, -1).sum(-1)
+
+    # E_q[log p(x|z)] via Monte Carlo average across S: [B]
+    E_log_px_z = log_px_z.mean(dim=1)
+
+    # KL(q||p): [B]
+    kl_q_p = kl_divergence_gaussian(mu, log_var)
+
+    # ELBO per sample and batch means (retain grads)
+    elbo_per_sample = E_log_px_z - beta * kl_q_p          # [B]
+    elbo = elbo_per_sample.mean()                         # scalar
+    E_log_px_z_mean = E_log_px_z.mean()                   # scalar
+    beta_kl_q_p_mean = beta * kl_q_p.mean()               # scalar
+
+    return ELBOComponents(
+        elbo=elbo,
+        log_likelihood=E_log_px_z_mean,
+        beta_kl_divergence=beta_kl_q_p_mean,
+    )

pyautoencoder-1.0.6/pyautoencoder/loss/wrapper.py

@@ -0,0 +1,213 @@
+"""Loss wrappers for easy loss computation and tracking."""
+from dataclasses import dataclass
+from typing import Dict, Optional, Union
+import math
+import torch
+
+from .base import LikelihoodType, log_likelihood
+from .vae import compute_ELBO
+from ..models.autoencoder import AEOutput
+from ..models.variational.vae import VAEOutput
+
+LN2 = math.log(2.0)
+LOG_2PI = math.log(2.0 * math.pi)  # for Gaussian σ²=1 diagnostics
+
+@dataclass
+class LossComponents:
+    """
+    Container for loss components with detailed metrics.
+
+    Args:
+        total (torch.Tensor): Scalar loss to optimize (already reduced over batch).
+        components (Dict[str, torch.Tensor]): Named scalar terms that compose the loss
+            (e.g., 'negative_log_likelihood', 'beta_kl_divergence').
+        metrics (Optional[Dict[str, torch.Tensor]]): Additional scalar diagnostics
+            (e.g., per-dimension metrics in nats/bits, KL per latent dimension).
+
+    Notes:
+        - All values are batch means unless specified otherwise.
+        - Metrics are intended for logging/monitoring and do not affect optimization directly.
+    """
+    total: torch.Tensor
+    components: Dict[str, torch.Tensor]
+    metrics: Optional[Dict[str, torch.Tensor]] = None
+
+class BaseLoss:
+    """Base class for all loss functions."""
+    def __call__(self, *args, **kwargs) -> LossComponents:
+        raise NotImplementedError
+
+class VAELoss(BaseLoss):
+    def __init__(
+        self,
+        beta: float = 1.0,
+        likelihood: Union[str, LikelihoodType] = LikelihoodType.GAUSSIAN,
+    ):
+        """
+        Loss function for Variational Autoencoders (β-VAE style).
+        The optimized loss is the *negative ELBO*. 
+        Uses negative log-likelihood (NLL) of reconstructions:
+        - Gaussian (σ²=1): per-dim NLL = 0.5·[(x − x_hat)² + log(2π)].
+        - Bernoulli (logits): per-dim NLL = BCEWithLogits(x_hat, x).
+
+        Args:
+            beta (float): Weighting factor for the KL term (β-VAE).
+            likelihood (Union[str, LikelihoodType]): Likelihood model for p(x|z).
+                Supported: 'gaussian' (σ²=1) or 'bernoulli' (logits).
+        """
+        self.beta = beta
+        self.likelihood = likelihood
+
+    def __call__(self, x: torch.Tensor, model_output: VAEOutput) -> LossComponents:
+        """
+        Computes VAE loss components and size-normalized diagnostics.
+
+        Args:
+            x (torch.Tensor): Ground truth inputs, shape [B, ...].
+            model_output (VAEOutput): from the VAE forward pass, dataclass with:
+                - x_hat (torch.Tensor): Reconstructed samples, shape [B, S, ...],
+                                        where S is the number of Monte Carlo samples from q(z|x).
+                - z (torch.Tensor): Latent samples (unused here).
+                - mu (torch.Tensor): Mean of q(z|x), shape [B, D_z].
+                - log_var (torch.Tensor): Log-variance of q(z|x), shape [B, D_z].
+
+        Returns:
+            LossComponents: Named container with:
+                - total (torch.Tensor): Scalar, negative ELBO (to minimize).
+                - components (Dict[str, torch.Tensor]):
+                    * 'negative_log_likelihood': Scalar, batch-mean -E_q[log p(x|z)] (nats).
+                    * 'beta_kl_divergence': Scalar, batch-mean β * KL(q||p) (nats).
+                - metrics (Dict[str, torch.Tensor]):
+                    * 'elbo': Scalar, batch-mean ELBO (nats).
+                    * 'nll_per_dim_nats': Scalar, -E_q[log p(x|z)] / D_x (nats/dim).
+                    * 'nll_per_dim_bits': Scalar, bits per dimension = nll_per_dim_nats / ln(2) (bits/dim).
+                    * 'beta_kl_per_latent_dim_nats': Scalar, β * KL / D_z (nats per latent dim).
+                    * 'beta_kl_per_latent_dim_bits': Scalar, beta_kl_per_latent_dim_nats / ln(2) (bits per latent dim).
+                    * 'mse_per_dim' (optional): Scalar, derived from Gaussian σ²=1 identity.
+
+        Notes:
+            - Reductions follow: sum over feature dimensions → mean over MC samples (if any) → mean over batch.
+            - log p(x|z) is computed by `log_likelihood`:
+                * Gaussian (σ²=1): per-dim NLL = 0.5·MSE + 0.5·log(2π).
+                * Bernoulli (logits): per-dim NLL = BCEWithLogits.
+            - For Gaussian (σ²=1), 'mse_per_dim' is computed via:
+                MSE_per_dim = 2·NLL_per_dim − log(2π), clamped to be ≥ 0.
+        """
+        x_hat = model_output.x_hat
+        mu = model_output.mu
+        log_var = model_output.log_var
+
+        elbo_components = compute_ELBO(
+            x=x,
+            x_hat=x_hat,
+            mu=mu,
+            log_var=log_var,
+            likelihood=self.likelihood,
+            beta=self.beta,
+        )
+
+        D_x = x[0].numel()
+        D_z = mu.size(-1)
+
+        # Per-dimension / per-latent-dimension metrics
+        nll_per_dim_nats = -elbo_components.log_likelihood / D_x          # nats/dim
+        nll_per_dim_bits = nll_per_dim_nats / LN2                         # bits/dim
+
+        beta_kl_per_latent_dim_nats = elbo_components.beta_kl_divergence / D_z      # nats/latent-dim
+        beta_kl_per_latent_dim_bits = beta_kl_per_latent_dim_nats / LN2             # bits/latent-dim
+
+        metrics: Dict[str, torch.Tensor] = {
+            'elbo': elbo_components.elbo.detach().cpu(),
+            'nll_per_dim_nats': nll_per_dim_nats.detach().cpu(),
+            'nll_per_dim_bits': nll_per_dim_bits.detach().cpu(),
+            'beta_kl_per_latent_dim_nats': beta_kl_per_latent_dim_nats.detach().cpu(),
+            'beta_kl_per_latent_dim_bits': beta_kl_per_latent_dim_bits.detach().cpu(),
+        }
+
+        # Extra: derive MSE/dim for Gaussian(σ²=1)
+        if self.likelihood == LikelihoodType.GAUSSIAN:
+            # NLL_per_dim = 0.5*MSE_per_dim + 0.5*log(2π) ⇒ MSE_per_dim = 2*NLL_per_dim − log(2π)
+            mse_per_dim = torch.clamp(2.0 * nll_per_dim_nats - LOG_2PI, min=0.0)
+            metrics['mse_per_dim'] = mse_per_dim.detach().cpu()
+
+        return LossComponents(
+            total=-elbo_components.elbo,  # minimize negative ELBO
+            components={
+                'negative_log_likelihood': -elbo_components.log_likelihood,
+                'beta_kl_divergence': elbo_components.beta_kl_divergence,
+            },
+            metrics=metrics,
+        )
+
+class AELoss(BaseLoss):
+    def __init__(self, likelihood: Union[str, LikelihoodType] = LikelihoodType.GAUSSIAN):
+        """
+        Loss function for standard Autoencoders.
+
+        Uses negative log-likelihood (NLL) of reconstructions:
+            - Gaussian (σ²=1): per-dim NLL = 0.5·[(x − x_hat)² + log(2π)].
+            - Bernoulli (logits): per-dim NLL = BCEWithLogits(x_hat, x).
+
+        Args:
+            likelihood (Union[str, LikelihoodType]): Likelihood model for p(x|z).
+                Supported: 'gaussian' (σ²=1) or 'bernoulli' (logits).
+        """
+        self.likelihood = likelihood
+
+    def __call__(self, x: torch.Tensor, model_output: AEOutput) -> LossComponents:
+        """
+        Computes Autoencoder reconstruction loss and size-normalized diagnostics.
+
+        Args:
+            x (torch.Tensor): Ground truth inputs, shape [B, ...].
+            model_output (AEOutput): from the AE forward pass, dataclass containing:
+                - x_hat (torch.Tensor): Reconstructions, shape [B, ...].
+                - z (torch.Tensor): Latent samples (unused here).
+
+        Returns:
+            LossComponents: Named container with:
+                - total (torch.Tensor): Scalar, batch-mean reconstruction loss (NLL in nats).
+                - components (Dict[str, torch.Tensor]):
+                    * 'negative_log_likelihood': Scalar, same as total.
+                - metrics (Dict[str, torch.Tensor]):
+                    * 'nll_per_dim_nats': Scalar, NLL / D_x (nats/dim).
+                    * 'nll_per_dim_bits': Scalar, bits per dimension = nll_per_dim_nats / ln(2) (bits/dim).
+                    * 'mse_per_dim' (optional): Scalar, derived for Gaussian σ²=1.
+
+        Notes:
+            - Reductions follow: elementwise log-likelihood → sum over feature dimensions
+              → mean over batch.
+            - For Gaussian (σ²=1), 'mse_per_dim' is computed via:
+                MSE_per_dim = 2·(NLL_per_dim) − log(2π), clamped to be ≥ 0.
+            - Ensure inputs match the likelihood’s expected scale:
+                * Gaussian: continuous data (typically standardized).
+                * Bernoulli: targets in [0, 1], predictions given as logits.
+        """
+        x_hat = model_output.x_hat
+
+        B = x.size(0)
+        D_x = x[0].numel()
+
+        # Elementwise log-likelihood → per-sample sum → batch mean
+        ll_elem = log_likelihood(x, x_hat, likelihood=self.likelihood)         # [B, ...]
+        ll_per_sample = ll_elem.reshape(B, -1).sum(-1)                         # [B]
+        nll = (-ll_per_sample).mean()                                          # scalar NLL
+
+        # Per-dim diagnostics
+        nll_per_dim_nats = nll / D_x                                           # nats/dim
+        nll_per_dim_bits = nll_per_dim_nats / LN2                              # bits/dim
+
+        metrics: Dict[str, torch.Tensor] = {
+            'nll_per_dim_nats': nll_per_dim_nats.detach().cpu(),
+            'nll_per_dim_bits': nll_per_dim_bits.detach().cpu(),
+        }
+
+        if self.likelihood == LikelihoodType.GAUSSIAN:
+            mse_per_dim = torch.clamp(2.0 * nll_per_dim_nats - LOG_2PI, min=0.0)
+            metrics['mse_per_dim'] = mse_per_dim.detach().cpu()
+
+        return LossComponents(
+            total=nll,
+            components={'negative_log_likelihood': nll},
+            metrics=metrics,
+        )