Models API

Auto-generated from docstrings.

ecgen.families.base

Shared base class for generative-model families.

Capabilities dataclass

What a family supports; checked when a config is built.

Attributes:

Name	Type	Description
conditional	bool	Accepts condition fields (a conditioning config).
cfg	bool	Supports classifier-free guidance at sampling time.
inpainting	bool	Can fill in unobserved parts of an ECG at sampling time.
likelihood	bool	Can evaluate (a bound on) log-likelihood.
stages	int	Number of training stages (2 for VQ-VAE + prior).

Source code in src/ecgen/families/base.py

@dataclass(frozen=True)
class Capabilities:
    """
    What a family supports; checked when a config is built.

    Attributes:
        conditional: Accepts condition fields (a ``conditioning`` config).
        cfg: Supports classifier-free guidance at sampling time.
        inpainting: Can fill in unobserved parts of an ECG at sampling time.
        likelihood: Can evaluate (a bound on) log-likelihood.
        stages: Number of training stages (2 for VQ-VAE + prior).
    """

    conditional: bool = False
    cfg: bool = False
    inpainting: bool = False
    likelihood: bool = False
    stages: int = 1

GenerativeModule

Bases: LightningModule

Base LightningModule for every family.

Each family owns its own training logic (loss, optimizers, manual vs. automatic optimization), so this class only fixes the shared contract:

• __init__(config=None, **kwargs) built with utils.config.resolve_config, so load_from_checkpoint(path) works from the flat saved hyperparameters;
• batches are read with :meth:get_signal, which accepts :class:~ecgen.types.ECGBatch and the legacy dict / tuple formats;
• conditional families build :attr:conditioner from their conditioning config with :meth:init_conditioning, and get conditions with :meth:condition (training) or :meth:prepare_condition (raw values at sampling time);
• :attr:data_info (lead names, fs, normalizers of the training data) is set by the runner and stored in checkpoints;
• :meth:sample returns [n, C, T] generated ECGs.

Class attributes

family: Family name. supports: :class:Capabilities. accepts: ConditionOutput slots the network consumes (subset of {"vector", "tokens", "dense"}).

Source code in src/ecgen/families/base.py

class GenerativeModule(pl.LightningModule):
    """
    Base ``LightningModule`` for every family.

    Each family owns its own training logic (loss, optimizers, manual vs.
    automatic optimization), so this class only fixes the shared contract:

    - ``__init__(config=None, **kwargs)`` built with ``utils.config.resolve_config``,
      so ``load_from_checkpoint(path)`` works from the flat saved hyperparameters;
    - batches are read with :meth:`get_signal`, which accepts :class:`~ecgen.types.ECGBatch`
      and the legacy dict / tuple formats;
    - conditional families build :attr:`conditioner` from their ``conditioning``
      config with :meth:`init_conditioning`, and get conditions with
      :meth:`condition` (training) or :meth:`prepare_condition` (raw values at
      sampling time);
    - :attr:`data_info` (lead names, fs, normalizers of the training data) is set
      by the runner and stored in checkpoints;
    - :meth:`sample` returns ``[n, C, T]`` generated ECGs.

    Class attributes:
        family: Family name.
        supports: :class:`Capabilities`.
        accepts: ConditionOutput slots the network consumes
            (subset of ``{"vector", "tokens", "dense"}``).
    """

    family: ClassVar[str] = ""
    supports: ClassVar[Capabilities] = Capabilities()
    accepts: ClassVar[FrozenSet[str]] = frozenset()

    conditioner: Optional[CompositeConditioner]

    def __init__(self) -> None:
        super().__init__()
        self.conditioner = None
        self.data_info: Dict[str, Any] = {}

    # ----------------------------------------------------------------- batch

    @staticmethod
    def get_signal(batch: Any) -> Tensor:
        return get_signal(batch)

    # ---------------------------------------------------------- conditioning

    def init_conditioning(self, spec: Optional[Mapping[str, Any]]) -> None:
        """Build :attr:`conditioner` from a ``conditioning`` config (``None`` = unconditional)."""
        conditioner = build_conditioner(spec)
        if conditioner is not None:
            if not self.supports.conditional:
                raise ValueError(f"{type(self).__name__} does not support conditioning")
            extra = conditioner.slots - set(self.accepts)
            if extra:
                raise ValueError(
                    f"{type(self).__name__} accepts condition slots {sorted(self.accepts)}, "
                    f"but the conditioning config produces {sorted(extra)}"
                )
        self.conditioner = conditioner

    @property
    def is_conditional(self) -> bool:
        return self.conditioner is not None

    def condition(self, batch: Any, drop: Drop = None) -> Optional[ConditionOutput]:
        """Encode a training batch's conditions (``None`` for unconditional models)."""
        if self.conditioner is None:
            return None
        if not isinstance(batch, ECGBatch):
            raise TypeError(
                "Conditional models need ECGBatch batches (use ECGDataModule)"
            )
        obs = batch.signal if batch.obs_mask is not None else None
        return self.conditioner(
            batch.cond,
            batch.cond_missing,
            obs=obs,
            obs_mask=batch.obs_mask,
            batch_size=len(batch),
            drop=drop,
        )

    def prepare_condition(
        self, cond: Optional[Mapping[str, Any]], n: int
    ) -> Tuple[Dict[str, Tensor], Dict[str, Tensor]]:
        """
        Turn raw condition values into model inputs for ``n`` samples.

        Numeric fields with a stored normalizer are given in raw units, either as
        a list/array of length F (broadcast to all samples) or ``[n, F]``; NaN
        or ``None`` marks a missing value. Other fields are passed as tensors (or
        anything ``torch.as_tensor`` accepts), broadcast if they have no batch
        dimension. Fields not given are left out (the model uses its null
        embedding for them).

        Returns:
            ``(cond, cond_missing)`` on this module's device.
        """
        device = self.device
        out: Dict[str, Tensor] = {}
        missing: Dict[str, Tensor] = {}
        normalizers = self.data_info.get("normalizers") or {}
        for field, value in (cond or {}).items():
            if field in normalizers:
                norm = FeatureNormalizer.from_dict(normalizers[field])
                raw = np.array(value, dtype=np.float64).reshape(-1, len(norm.columns))
                raw = np.broadcast_to(raw, (n, len(norm.columns)))
                vals, miss = zip(*(norm.transform(row) for row in raw))
                out[field] = torch.as_tensor(np.stack(vals), device=device)
                missing[field] = torch.as_tensor(np.stack(miss), device=device)
            elif isinstance(value, str) or (
                isinstance(value, (list, tuple)) and value and isinstance(value[0], str)
            ):
                texts = [value] * n if isinstance(value, str) else list(value)
                if len(texts) != n:
                    raise ValueError(
                        f"{field}: got {len(texts)} strings for {n} samples"
                    )
                out[field] = texts
            else:
                t = torch.as_tensor(value, device=device)
                if t.dim() == 0 or (t.shape[0] != n):
                    t = t.unsqueeze(0).expand(n, *t.shape)
                out[field] = t
        return out, missing

    def encode_condition(
        self,
        n: int,
        cond: Optional[Mapping[str, Any]] = None,
        obs: Optional[Tensor] = None,
        obs_mask: Optional[Tensor] = None,
        drop: Drop = None,
    ) -> Optional[ConditionOutput]:
        """Raw sampling-time conditions → :class:`ConditionOutput` (``None`` if unconditional)."""
        if self.conditioner is None:
            if cond or obs is not None:
                raise ValueError(
                    f"{type(self).__name__} was trained without conditioning"
                )
            return None
        values, missing = self.prepare_condition(cond, n)
        if obs is not None:
            obs = obs.to(self.device)
            obs_mask = (
                torch.ones_like(obs)
                if obs_mask is None
                else obs_mask.to(self.device, obs.dtype)
            )
        return self.conditioner(
            values, missing, obs=obs, obs_mask=obs_mask, batch_size=n, drop=drop
        )

    # ------------------------------------------------------------- data info

    def set_data_info(self, info: Mapping[str, Any]) -> None:
        """Record the training data contract (lead names, fs, seq_length, normalizers)."""
        self.data_info = dict(info)

    def on_save_checkpoint(self, checkpoint: Dict[str, Any]) -> None:
        from ecgen.registry import FAMILIES

        checkpoint["ecgen_data_info"] = self.data_info
        checkpoint["ecgen_family"] = FAMILIES.name_of(self)

    def on_load_checkpoint(self, checkpoint: Dict[str, Any]) -> None:
        self.data_info = dict(checkpoint.get("ecgen_data_info") or {})

    @property
    def fixed_seq_length(self) -> Optional[int]:
        """Signal length the network is built for, or ``None`` if any length works."""
        return None

    # -------------------------------------------------------------- sampling

    def sample(self, n_samples: int = 16, **kwargs: Any) -> Tensor:
        raise NotImplementedError(f"{type(self).__name__} does not implement sample()")

fixed_seq_length property

fixed_seq_length: Optional[int]

Signal length the network is built for, or None if any length works.

init_conditioning

init_conditioning(spec: Optional[Mapping[str, Any]]) -> None

Build :attr:conditioner from a conditioning config (None = unconditional).

Source code in src/ecgen/families/base.py

def init_conditioning(self, spec: Optional[Mapping[str, Any]]) -> None:
    """Build :attr:`conditioner` from a ``conditioning`` config (``None`` = unconditional)."""
    conditioner = build_conditioner(spec)
    if conditioner is not None:
        if not self.supports.conditional:
            raise ValueError(f"{type(self).__name__} does not support conditioning")
        extra = conditioner.slots - set(self.accepts)
        if extra:
            raise ValueError(
                f"{type(self).__name__} accepts condition slots {sorted(self.accepts)}, "
                f"but the conditioning config produces {sorted(extra)}"
            )
    self.conditioner = conditioner

condition

condition(batch: Any, drop: Drop = None) -> Optional[ConditionOutput]

Encode a training batch's conditions (None for unconditional models).

Source code in src/ecgen/families/base.py

def condition(self, batch: Any, drop: Drop = None) -> Optional[ConditionOutput]:
    """Encode a training batch's conditions (``None`` for unconditional models)."""
    if self.conditioner is None:
        return None
    if not isinstance(batch, ECGBatch):
        raise TypeError(
            "Conditional models need ECGBatch batches (use ECGDataModule)"
        )
    obs = batch.signal if batch.obs_mask is not None else None
    return self.conditioner(
        batch.cond,
        batch.cond_missing,
        obs=obs,
        obs_mask=batch.obs_mask,
        batch_size=len(batch),
        drop=drop,
    )

prepare_condition

prepare_condition(cond: Optional[Mapping[str, Any]], n: int) -> Tuple[Dict[str, Tensor], Dict[str, Tensor]]

Turn raw condition values into model inputs for n samples.

Numeric fields with a stored normalizer are given in raw units, either as a list/array of length F (broadcast to all samples) or [n, F]; NaN or None marks a missing value. Other fields are passed as tensors (or anything torch.as_tensor accepts), broadcast if they have no batch dimension. Fields not given are left out (the model uses its null embedding for them).

Returns:

Type	Description
Tuple[Dict[str, Tensor], Dict[str, Tensor]]	(cond, cond_missing) on this module's device.

Source code in src/ecgen/families/base.py

def prepare_condition(
    self, cond: Optional[Mapping[str, Any]], n: int
) -> Tuple[Dict[str, Tensor], Dict[str, Tensor]]:
    """
    Turn raw condition values into model inputs for ``n`` samples.

    Numeric fields with a stored normalizer are given in raw units, either as
    a list/array of length F (broadcast to all samples) or ``[n, F]``; NaN
    or ``None`` marks a missing value. Other fields are passed as tensors (or
    anything ``torch.as_tensor`` accepts), broadcast if they have no batch
    dimension. Fields not given are left out (the model uses its null
    embedding for them).

    Returns:
        ``(cond, cond_missing)`` on this module's device.
    """
    device = self.device
    out: Dict[str, Tensor] = {}
    missing: Dict[str, Tensor] = {}
    normalizers = self.data_info.get("normalizers") or {}
    for field, value in (cond or {}).items():
        if field in normalizers:
            norm = FeatureNormalizer.from_dict(normalizers[field])
            raw = np.array(value, dtype=np.float64).reshape(-1, len(norm.columns))
            raw = np.broadcast_to(raw, (n, len(norm.columns)))
            vals, miss = zip(*(norm.transform(row) for row in raw))
            out[field] = torch.as_tensor(np.stack(vals), device=device)
            missing[field] = torch.as_tensor(np.stack(miss), device=device)
        elif isinstance(value, str) or (
            isinstance(value, (list, tuple)) and value and isinstance(value[0], str)
        ):
            texts = [value] * n if isinstance(value, str) else list(value)
            if len(texts) != n:
                raise ValueError(
                    f"{field}: got {len(texts)} strings for {n} samples"
                )
            out[field] = texts
        else:
            t = torch.as_tensor(value, device=device)
            if t.dim() == 0 or (t.shape[0] != n):
                t = t.unsqueeze(0).expand(n, *t.shape)
            out[field] = t
    return out, missing

encode_condition

encode_condition(n: int, cond: Optional[Mapping[str, Any]] = None, obs: Optional[Tensor] = None, obs_mask: Optional[Tensor] = None, drop: Drop = None) -> Optional[ConditionOutput]

Raw sampling-time conditions → :class:ConditionOutput (None if unconditional).

Source code in src/ecgen/families/base.py

def encode_condition(
    self,
    n: int,
    cond: Optional[Mapping[str, Any]] = None,
    obs: Optional[Tensor] = None,
    obs_mask: Optional[Tensor] = None,
    drop: Drop = None,
) -> Optional[ConditionOutput]:
    """Raw sampling-time conditions → :class:`ConditionOutput` (``None`` if unconditional)."""
    if self.conditioner is None:
        if cond or obs is not None:
            raise ValueError(
                f"{type(self).__name__} was trained without conditioning"
            )
        return None
    values, missing = self.prepare_condition(cond, n)
    if obs is not None:
        obs = obs.to(self.device)
        obs_mask = (
            torch.ones_like(obs)
            if obs_mask is None
            else obs_mask.to(self.device, obs.dtype)
        )
    return self.conditioner(
        values, missing, obs=obs, obs_mask=obs_mask, batch_size=n, drop=drop
    )

set_data_info

set_data_info(info: Mapping[str, Any]) -> None

Record the training data contract (lead names, fs, seq_length, normalizers).

Source code in src/ecgen/families/base.py

def set_data_info(self, info: Mapping[str, Any]) -> None:
    """Record the training data contract (lead names, fs, seq_length, normalizers)."""
    self.data_info = dict(info)

ecgen.nn.blocks

ResidualBlock1D

Bases: Module

1D Residual block with group normalization

Source code in src/ecgen/nn/blocks.py

class ResidualBlock1D(nn.Module):
    """1D Residual block with group normalization"""

    def __init__(self, in_channels: int, out_channels: int) -> None:
        super().__init__()
        self.conv1 = nn.Conv1d(in_channels, out_channels, kernel_size=3, padding=1)
        self.conv2 = nn.Conv1d(out_channels, out_channels, kernel_size=3, padding=1)
        self.norm1 = nn.GroupNorm(8, out_channels)
        self.norm2 = nn.GroupNorm(8, out_channels)

        if in_channels != out_channels:
            self.shortcut = nn.Conv1d(in_channels, out_channels, kernel_size=1)
        else:
            self.shortcut = nn.Identity()

    def forward(self, x: Tensor) -> Tensor:
        residual = self.shortcut(x)

        h = self.conv1(x)
        h = self.norm1(h)
        h = F.silu(h)

        h = self.conv2(h)
        h = self.norm2(h)
        h = F.silu(h)

        return h + residual

ecgen.nn.encoders

Encoder1D

Bases: Module

1D Encoder for ECG signals

Source code in src/ecgen/nn/encoders.py

class Encoder1D(nn.Module):
    """1D Encoder for ECG signals"""

    def __init__(
        self,
        in_channels: int = 12,
        base_channels: int = 64,
        latent_channels: int = 8,
        channel_multipliers: tuple[int, ...] = (1, 2, 4, 4),
        num_res_blocks: int = 2,
        cond_dim: int = 0,
        extra_in_channels: int = 0,
    ) -> None:
        super().__init__()

        self.in_channels = in_channels
        self.latent_channels = latent_channels

        # extra_in_channels: dense conditions concatenated to the input;
        # cond_dim > 0: FiLM after every residual block (both default off, so
        # unconditional models keep their original parameters)
        self.conv_in = nn.Conv1d(
            in_channels + extra_in_channels, base_channels, kernel_size=7, padding=3
        )

        self.down_blocks = nn.ModuleList()
        ch = base_channels

        for i, mult in enumerate(channel_multipliers):
            out_ch = base_channels * mult

            for _ in range(num_res_blocks):
                self.down_blocks.append(ResidualBlock1D(ch, out_ch))
                ch = out_ch

            if i < len(channel_multipliers) - 1:
                self.down_blocks.append(
                    nn.Conv1d(ch, ch, kernel_size=4, stride=2, padding=1)
                )

        self.mid_block1 = ResidualBlock1D(ch, ch)
        self.mid_block2 = ResidualBlock1D(ch, ch)

        self.norm_out = nn.GroupNorm(8, ch)
        self.conv_out = nn.Conv1d(ch, latent_channels * 2, kernel_size=3, padding=1)

        self.films = nn.ModuleDict()
        if cond_dim > 0:
            for i, block in enumerate(self.down_blocks):
                if isinstance(block, ResidualBlock1D):
                    self.films[str(i)] = FiLM(cond_dim, block.conv2.out_channels)
            self.films["mid"] = FiLM(cond_dim, ch)

    def forward(
        self, x: Tensor, c: Optional[Tensor] = None, dense: Optional[Tensor] = None
    ) -> tuple[Tensor, Tensor]:
        if dense is not None:
            x = torch.cat([x, dense], dim=1)
        h = self.conv_in(x)

        for i, block in enumerate(self.down_blocks):
            h = block(h)
            if str(i) in self.films:
                h = self.films[str(i)](h, c)

        h = self.mid_block1(h)
        h = self.mid_block2(h)
        if "mid" in self.films:
            h = self.films["mid"](h, c)

        h = self.norm_out(h)
        h = F.silu(h)
        h = self.conv_out(h)

        mean, logvar = torch.chunk(h, 2, dim=1)

        return mean, logvar

ecgen.nn.decoders

Decoder1D

Bases: Module

1D Decoder for ECG signals

Source code in src/ecgen/nn/decoders.py

class Decoder1D(nn.Module):
    """1D Decoder for ECG signals"""

    def __init__(
        self,
        out_channels: int = 12,
        base_channels: int = 64,
        latent_channels: int = 8,
        channel_multipliers: tuple[int, ...] = (1, 2, 4, 4),
        num_res_blocks: int = 2,
        cond_dim: int = 0,
        extra_in_channels: int = 0,
    ) -> None:
        super().__init__()

        self.out_channels = out_channels
        self.latent_channels = latent_channels

        channel_multipliers = tuple(reversed(channel_multipliers))

        ch = base_channels * channel_multipliers[0]
        # extra_in_channels: dense conditions (at latent resolution) concatenated
        # to z; cond_dim > 0: FiLM after the mid blocks and every residual block
        self.conv_in = nn.Conv1d(
            latent_channels + extra_in_channels, ch, kernel_size=3, padding=1
        )

        self.mid_block1 = ResidualBlock1D(ch, ch)
        self.mid_block2 = ResidualBlock1D(ch, ch)

        self.up_blocks = nn.ModuleList()

        for i, mult in enumerate(channel_multipliers):
            out_ch = base_channels * mult

            for _ in range(num_res_blocks):
                self.up_blocks.append(ResidualBlock1D(ch, out_ch))
                ch = out_ch

            if i < len(channel_multipliers) - 1:
                self.up_blocks.append(
                    nn.ConvTranspose1d(ch, ch, kernel_size=4, stride=2, padding=1)
                )

        self.norm_out = nn.GroupNorm(8, ch)
        self.conv_out = nn.Conv1d(ch, out_channels, kernel_size=7, padding=3)

        self.films = nn.ModuleDict()
        if cond_dim > 0:
            self.films["mid"] = FiLM(cond_dim, base_channels * channel_multipliers[0])
            for i, block in enumerate(self.up_blocks):
                if isinstance(block, ResidualBlock1D):
                    self.films[str(i)] = FiLM(cond_dim, block.conv2.out_channels)

    def forward(
        self, z: Tensor, c: Optional[Tensor] = None, dense: Optional[Tensor] = None
    ) -> Tensor:
        if dense is not None:
            z = torch.cat([z, dense], dim=1)
        h = self.conv_in(z)

        h = self.mid_block1(h)
        h = self.mid_block2(h)
        if "mid" in self.films:
            h = self.films["mid"](h, c)

        for i, block in enumerate(self.up_blocks):
            h = block(h)
            if str(i) in self.films:
                h = self.films[str(i)](h, c)

        h = self.norm_out(h)
        h = F.silu(h)
        h = self.conv_out(h)

        return h

ecgen.families.vae.config

ecgen.families.vae.vae_1d

GlobalPrior

Bases: Module

Learned p(z | c) for a global latent [B, D]; zero-initialised (starts as N(0, I)).

Source code in src/ecgen/families/vae/vae_1d.py

class GlobalPrior(nn.Module):
    """Learned ``p(z | c)`` for a global latent ``[B, D]``; zero-initialised (starts as N(0, I))."""

    def __init__(self, latent_dim: int, cond_dim: int) -> None:
        super().__init__()
        self.proj = nn.Linear(cond_dim, 2 * latent_dim)
        nn.init.zeros_(self.proj.weight)
        nn.init.zeros_(self.proj.bias)
        self.latent_dim = latent_dim

    def forward(
        self, c: Optional[Tensor], batch_size: int, device
    ) -> tuple[Tensor, Tensor]:
        if c is None:
            zeros = torch.zeros(batch_size, self.latent_dim, device=device)
            return zeros, zeros
        mean, logvar = self.proj(c).chunk(2, dim=-1)
        return mean, logvar

ConditionalPrior

Bases: Module

Learned prior p(z | c): a Gaussian per latent position.

Global conditions give a per-channel mean/log-variance; dense conditions (at latent resolution) add a per-position term. Zero-initialised, so it starts as N(0, I).

Source code in src/ecgen/families/vae/vae_1d.py

class ConditionalPrior(nn.Module):
    """
    Learned prior ``p(z | c)``: a Gaussian per latent position.

    Global conditions give a per-channel mean/log-variance; dense conditions
    (at latent resolution) add a per-position term. Zero-initialised, so it
    starts as N(0, I).
    """

    def __init__(
        self, latent_channels: int, cond_dim: int = 0, dense_channels: int = 0
    ) -> None:
        super().__init__()
        self.vector = nn.Linear(cond_dim, 2 * latent_channels) if cond_dim else None
        self.dense = (
            nn.Conv1d(dense_channels, 2 * latent_channels, kernel_size=3, padding=1)
            if dense_channels
            else None
        )
        for m in (self.vector, self.dense):
            if m is not None:
                nn.init.zeros_(m.weight)
                nn.init.zeros_(m.bias)
        self.latent_channels = latent_channels

    def forward(
        self,
        length: int,
        c: Optional[Tensor],
        dense: Optional[Tensor],
        batch_size: int,
        device,
    ) -> tuple[Tensor, Tensor]:
        params = torch.zeros(
            batch_size, 2 * self.latent_channels, length, device=device
        )
        if self.vector is not None and c is not None:
            params = params + self.vector(c)[..., None]
        if self.dense is not None and dense is not None:
            params = params + self.dense(dense)
        mean, logvar = params.chunk(2, dim=1)
        return mean, logvar

VAE1D

Bases: Module

Variational Autoencoder for ECG signals; a conditional VAE when cond_dim (global conditions, via FiLM) or dense_channels (time-aligned conditions, concatenated) is set.

latent_type="spatial" (default) keeps a [latent_channels, T / f] latent map with a per-position Gaussian prior — good reconstructions and a first stage for latent priors, but samples straight from the prior are not coherent ECGs. latent_type="global" uses one [global_latent_dim] vector per ECG (flattened encoder output -> linear; linear -> decoder input map), whose prior samples decode to coherent ECGs; it needs seq_length.

Source code in src/ecgen/families/vae/vae_1d.py

class VAE1D(nn.Module):
    """
    Variational Autoencoder for ECG signals; a conditional VAE when
    ``cond_dim`` (global conditions, via FiLM) or ``dense_channels``
    (time-aligned conditions, concatenated) is set.

    ``latent_type="spatial"`` (default) keeps a ``[latent_channels, T / f]``
    latent map with a per-position Gaussian prior — good reconstructions and a
    first stage for latent priors, but samples straight from the prior are not
    coherent ECGs. ``latent_type="global"`` uses one ``[global_latent_dim]``
    vector per ECG (flattened encoder output -> linear; linear -> decoder input
    map), whose prior samples decode to coherent ECGs; it needs ``seq_length``.
    """

    def __init__(
        self,
        in_channels: int = 12,
        base_channels: int = 64,
        latent_channels: int = 8,
        channel_multipliers: tuple[int, ...] = (1, 2, 4, 4),
        num_res_blocks: int = 2,
        cond_dim: int = 0,
        dense_channels: int = 0,
        learned_prior: bool = False,
        latent_type: str = "spatial",
        global_latent_dim: int = 64,
        seq_length: Optional[int] = None,
    ) -> None:
        super().__init__()
        self.downsample_factor = 2 ** (len(channel_multipliers) - 1)
        if latent_type not in ("spatial", "global"):
            raise ValueError(
                f"latent_type must be 'spatial' or 'global', got {latent_type!r}"
            )
        self.latent_type = latent_type
        self.latent_channels = latent_channels

        self.encoder = Encoder1D(
            in_channels=in_channels,
            base_channels=base_channels,
            latent_channels=latent_channels,
            channel_multipliers=channel_multipliers,
            num_res_blocks=num_res_blocks,
            cond_dim=cond_dim,
            extra_in_channels=dense_channels,
        )

        self.decoder = Decoder1D(
            out_channels=in_channels,
            base_channels=base_channels,
            latent_channels=latent_channels,
            channel_multipliers=channel_multipliers,
            num_res_blocks=num_res_blocks,
            cond_dim=cond_dim,
            extra_in_channels=dense_channels,
        )

        self.prior: Optional[nn.Module] = None
        if latent_type == "global":
            if not seq_length or seq_length % self.downsample_factor:
                raise ValueError(
                    "latent_type='global' needs seq_length divisible by the downsampling"
                )
            self.map_length = seq_length // self.downsample_factor
            flat = 2 * latent_channels * self.map_length
            self.to_global = nn.Linear(flat, 2 * global_latent_dim)
            self.from_global = nn.Linear(
                global_latent_dim, latent_channels * self.map_length
            )
            self.global_latent_dim = global_latent_dim
            if learned_prior:
                self.prior = (
                    GlobalPrior(global_latent_dim, cond_dim) if cond_dim else None
                )
                if dense_channels:
                    raise ValueError(
                        "learned_prior with latent_type='global' supports vector conditions only"
                    )
        elif learned_prior:
            self.prior = ConditionalPrior(latent_channels, cond_dim, dense_channels)

    def latent_dense(self, dense: Optional[Tensor], length: int) -> Optional[Tensor]:
        """Dense conditions pooled to the latent length."""
        return None if dense is None else F.adaptive_avg_pool1d(dense, length)

    def encode(
        self, x: Tensor, c: Optional[Tensor] = None, dense: Optional[Tensor] = None
    ) -> tuple[Tensor, Tensor]:
        """Encode input to latent distribution"""
        mean, logvar = self.encoder(x, c, dense)
        if self.latent_type == "global":
            mean, logvar = self.to_global(
                torch.cat([mean, logvar], dim=1).flatten(1)
            ).chunk(2, dim=-1)
        return mean, logvar

    def reparameterize(self, mean: Tensor, logvar: Tensor) -> Tensor:
        """Reparameterization trick"""
        std = torch.exp(0.5 * logvar)
        eps = torch.randn_like(std)
        return mean + eps * std

    def decode(
        self, z: Tensor, c: Optional[Tensor] = None, dense: Optional[Tensor] = None
    ) -> Tensor:
        """Decode latent to reconstruction (``dense`` at full resolution)."""
        if self.latent_type == "global":
            z = self.from_global(z).view(
                z.shape[0], self.latent_channels, self.map_length
            )
        return self.decoder(z, c, self.latent_dense(dense, z.shape[-1]))

    def prior_params(
        self,
        length: int,
        c: Optional[Tensor],
        dense: Optional[Tensor],
        batch_size: int,
        device,
    ) -> tuple[Tensor, Tensor]:
        """Mean and log-variance of ``p(z | c)`` (zeros for the standard normal prior).

        ``length`` is the latent map length (ignored for a global latent).
        """
        if self.latent_type == "global":
            if self.prior is None:
                zeros = torch.zeros(batch_size, self.global_latent_dim, device=device)
                return zeros, zeros
            return self.prior(c, batch_size, device)
        if self.prior is None:
            zeros = torch.zeros(
                batch_size, self.decoder.latent_channels, length, device=device
            )
            return zeros, zeros
        return self.prior(
            length, c, self.latent_dense(dense, length), batch_size, device
        )

    def latent_shape(self, seq_length: int) -> tuple[int, ...]:
        """Shape of one sample's latent."""
        if self.latent_type == "global":
            return (self.global_latent_dim,)
        return (self.latent_channels, seq_length // self.downsample_factor)

    def forward(
        self, x: Tensor, c: Optional[Tensor] = None, dense: Optional[Tensor] = None
    ) -> tuple[Tensor, Tensor, Tensor]:
        """Full forward pass"""
        mean, logvar = self.encode(x, c, dense)
        z = self.reparameterize(mean, logvar)
        recon = self.decode(z, c, dense)
        return recon, mean, logvar

    @torch.no_grad()
    def encode_to_latent(self, x: Tensor) -> Tensor:
        """Encode to latent (for diffusion training)"""
        mean, _ = self.encode(x)
        return mean

    @torch.no_grad()
    def decode_from_latent(self, z: Tensor) -> Tensor:
        """Decode from latent (for diffusion sampling)"""
        return self.decode(z)

latent_dense

latent_dense(dense: Optional[Tensor], length: int) -> Optional[Tensor]

Dense conditions pooled to the latent length.

Source code in src/ecgen/families/vae/vae_1d.py

def latent_dense(self, dense: Optional[Tensor], length: int) -> Optional[Tensor]:
    """Dense conditions pooled to the latent length."""
    return None if dense is None else F.adaptive_avg_pool1d(dense, length)

encode

encode(x: Tensor, c: Optional[Tensor] = None, dense: Optional[Tensor] = None) -> tuple[Tensor, Tensor]

Encode input to latent distribution

Source code in src/ecgen/families/vae/vae_1d.py

def encode(
    self, x: Tensor, c: Optional[Tensor] = None, dense: Optional[Tensor] = None
) -> tuple[Tensor, Tensor]:
    """Encode input to latent distribution"""
    mean, logvar = self.encoder(x, c, dense)
    if self.latent_type == "global":
        mean, logvar = self.to_global(
            torch.cat([mean, logvar], dim=1).flatten(1)
        ).chunk(2, dim=-1)
    return mean, logvar

reparameterize

reparameterize(mean: Tensor, logvar: Tensor) -> Tensor

Reparameterization trick

Source code in src/ecgen/families/vae/vae_1d.py

def reparameterize(self, mean: Tensor, logvar: Tensor) -> Tensor:
    """Reparameterization trick"""
    std = torch.exp(0.5 * logvar)
    eps = torch.randn_like(std)
    return mean + eps * std

decode

decode(z: Tensor, c: Optional[Tensor] = None, dense: Optional[Tensor] = None) -> Tensor

Decode latent to reconstruction (dense at full resolution).

Source code in src/ecgen/families/vae/vae_1d.py

def decode(
    self, z: Tensor, c: Optional[Tensor] = None, dense: Optional[Tensor] = None
) -> Tensor:
    """Decode latent to reconstruction (``dense`` at full resolution)."""
    if self.latent_type == "global":
        z = self.from_global(z).view(
            z.shape[0], self.latent_channels, self.map_length
        )
    return self.decoder(z, c, self.latent_dense(dense, z.shape[-1]))

prior_params

prior_params(length: int, c: Optional[Tensor], dense: Optional[Tensor], batch_size: int, device) -> tuple[Tensor, Tensor]

Mean and log-variance of p(z | c) (zeros for the standard normal prior).

length is the latent map length (ignored for a global latent).

Source code in src/ecgen/families/vae/vae_1d.py

def prior_params(
    self,
    length: int,
    c: Optional[Tensor],
    dense: Optional[Tensor],
    batch_size: int,
    device,
) -> tuple[Tensor, Tensor]:
    """Mean and log-variance of ``p(z | c)`` (zeros for the standard normal prior).

    ``length`` is the latent map length (ignored for a global latent).
    """
    if self.latent_type == "global":
        if self.prior is None:
            zeros = torch.zeros(batch_size, self.global_latent_dim, device=device)
            return zeros, zeros
        return self.prior(c, batch_size, device)
    if self.prior is None:
        zeros = torch.zeros(
            batch_size, self.decoder.latent_channels, length, device=device
        )
        return zeros, zeros
    return self.prior(
        length, c, self.latent_dense(dense, length), batch_size, device
    )

latent_shape

latent_shape(seq_length: int) -> tuple[int, ...]

Shape of one sample's latent.

Source code in src/ecgen/families/vae/vae_1d.py

def latent_shape(self, seq_length: int) -> tuple[int, ...]:
    """Shape of one sample's latent."""
    if self.latent_type == "global":
        return (self.global_latent_dim,)
    return (self.latent_channels, seq_length // self.downsample_factor)

forward

forward(x: Tensor, c: Optional[Tensor] = None, dense: Optional[Tensor] = None) -> tuple[Tensor, Tensor, Tensor]

Full forward pass

Source code in src/ecgen/families/vae/vae_1d.py

def forward(
    self, x: Tensor, c: Optional[Tensor] = None, dense: Optional[Tensor] = None
) -> tuple[Tensor, Tensor, Tensor]:
    """Full forward pass"""
    mean, logvar = self.encode(x, c, dense)
    z = self.reparameterize(mean, logvar)
    recon = self.decode(z, c, dense)
    return recon, mean, logvar

encode_to_latent

encode_to_latent(x: Tensor) -> Tensor

Encode to latent (for diffusion training)

Source code in src/ecgen/families/vae/vae_1d.py

@torch.no_grad()
def encode_to_latent(self, x: Tensor) -> Tensor:
    """Encode to latent (for diffusion training)"""
    mean, _ = self.encode(x)
    return mean

decode_from_latent

decode_from_latent(z: Tensor) -> Tensor

Decode from latent (for diffusion sampling)

Source code in src/ecgen/families/vae/vae_1d.py

@torch.no_grad()
def decode_from_latent(self, z: Tensor) -> Tensor:
    """Decode from latent (for diffusion sampling)"""
    return self.decode(z)

ecgen.families.vae.vae_lightning

VAELightning

Bases: GenerativeModule

PyTorch Lightning wrapper for VAE1D model.

With config.conditioning it is a conditional VAE: global conditions (vector) modulate the encoder and decoder with FiLM, time-aligned conditions (dense, e.g. a masked ECG) are concatenated to their inputs, and learned_prior replaces N(0, I) with p(z | c).

Source code in src/ecgen/families/vae/vae_lightning.py

@FAMILIES.register("vae")
class VAELightning(GenerativeModule):
    """
    PyTorch Lightning wrapper for VAE1D model.

    With ``config.conditioning`` it is a conditional VAE: global conditions
    (``vector``) modulate the encoder and decoder with FiLM, time-aligned
    conditions (``dense``, e.g. a masked ECG) are concatenated to their inputs,
    and ``learned_prior`` replaces N(0, I) with ``p(z | c)``.
    """

    family = "vae"
    supports = Capabilities(conditional=True)
    accepts = frozenset({"vector", "dense"})

    def __init__(
        self, config: VAEConfig | dict[str, Any] | None = None, **kwargs: Any
    ) -> None:
        super().__init__()
        config = resolve_config(VAEConfig, config, kwargs)
        self.save_hyperparameters(config.__dict__)
        self.config = config
        if config.recon_on not in ("all", "unobserved"):
            raise ValueError(
                f"recon_on must be 'all' or 'unobserved', got {config.recon_on!r}"
            )

        self.init_conditioning(config.conditioning)
        cond_dim = dense_channels = 0
        if self.conditioner is not None:
            if "vector" in self.conditioner.slots:
                cond_dim = self.conditioner.vector_dim
            dense_channels = self.conditioner.dense_channels

        self.vae = VAE1D(
            in_channels=config.in_channels,
            base_channels=config.base_channels,
            latent_channels=config.latent_channels,
            channel_multipliers=config.channel_multipliers,
            num_res_blocks=config.num_res_blocks,
            cond_dim=cond_dim,
            dense_channels=dense_channels,
            learned_prior=config.learned_prior,
            latent_type=config.latent_type,
            global_latent_dim=config.global_latent_dim,
            seq_length=config.seq_length,
        )
        if config.learned_prior and self.conditioner is None:
            raise ValueError("learned_prior needs a conditioning config")

        self._val_real_sample: Tensor | None = None
        self._val_recon_sample: Tensor | None = None

    def forward(self, x: Tensor) -> tuple[Tensor, Tensor, Tensor]:
        return self.vae(x)

    @property
    def fixed_seq_length(self) -> Optional[int]:
        return self.config.seq_length if self.config.latent_type == "global" else None

    def _step(self, batch: Any) -> tuple[Tensor, Tensor, Tensor, Tensor, Tensor]:
        ecgs = self.get_signal(batch)
        cond = self.condition(batch)
        if cond is None:
            recon, mean, logvar = self.vae(ecgs)
            total, rec, kl = vae_loss(recon, ecgs, mean, logvar, self.config.kl_weight)
            return total, rec, kl, ecgs, recon

        recon, mean, logvar = self.vae(ecgs, cond.vector, cond.dense)
        p_mean, p_logvar = self.vae.prior_params(
            ecgs.shape[-1] // self.vae.downsample_factor,
            cond.vector,
            cond.dense,
            ecgs.shape[0],
            ecgs.device,
        )
        mask = None
        if self.config.recon_on == "unobserved" and batch.obs_mask is not None:
            mask = 1.0 - batch.obs_mask
        total, rec, kl = cvae_loss(
            recon, ecgs, mean, logvar, p_mean, p_logvar, self.config.kl_weight, mask
        )
        return total, rec, kl, ecgs, recon

    def training_step(self, batch: Any, batch_idx: int) -> Tensor:
        total_loss, recon_loss, kl_loss, _, _ = self._step(batch)

        self.log(
            "train/total_loss", total_loss, on_step=True, on_epoch=True, prog_bar=True
        )
        self.log(
            "train/recon_loss", recon_loss, on_step=True, on_epoch=True, prog_bar=True
        )
        self.log("train/kl_loss", kl_loss, on_step=True, on_epoch=True, prog_bar=False)

        return total_loss

    def validation_step(self, batch: Any, batch_idx: int) -> Tensor:
        total_loss, recon_loss, kl_loss, ecgs, recon = self._step(batch)

        self.log(
            "val/total_loss", total_loss, on_step=False, on_epoch=True, prog_bar=True
        )
        self.log(
            "val/recon_loss", recon_loss, on_step=False, on_epoch=True, prog_bar=True
        )
        self.log("val/kl_loss", kl_loss, on_step=False, on_epoch=True, prog_bar=False)
        self.log("val_loss", total_loss, on_step=False, on_epoch=True, prog_bar=True)

        if batch_idx == 0:
            self._val_real_sample = ecgs[0].detach().cpu()
            self._val_recon_sample = recon[0].detach().cpu()

        return total_loss

    @torch.no_grad()
    def reconstruct(self, x: Tensor) -> Tensor:
        """Reconstruct input ECG signals (unconditional models)."""
        self.vae.eval()
        recon, _, _ = self.vae(x)
        return recon

    @torch.no_grad()
    def sample(
        self,
        n_samples: int = 16,
        seq_length: int = 5000,
        cond: Optional[Mapping[str, Any]] = None,
        obs: Optional[Tensor] = None,
        obs_mask: Optional[Tensor] = None,
        paste_observed: bool = True,
        **kwargs: Any,
    ) -> Tensor:
        """
        Sample ECGs: ``z`` from the prior (``p(z | c)`` if learned), then decode.

        Args:
            n_samples: Number of ECGs.
            seq_length: Length (multiple of the downsampling factor).
            cond: Raw condition values (see :meth:`prepare_condition`).
            obs: Partially observed ECG ``[n, C, T]``.
            obs_mask: Mask of ``obs`` (1 = observed).
            paste_observed: Copy the observed part of ``obs`` into the output.
        """
        self.vae.eval()
        device = next(self.vae.parameters()).device
        if self.config.latent_type == "global":
            seq_length = self.config.seq_length
        latent_length = seq_length // self.vae.downsample_factor
        c: Optional[ConditionOutput] = self.encode_condition(
            n_samples, cond, obs, obs_mask
        )
        vector = c.vector if c is not None else None
        dense = c.dense if c is not None else None

        mean, logvar = self.vae.prior_params(
            latent_length, vector, dense, n_samples, device
        )
        noise = torch.randn(
            n_samples, *self.vae.latent_shape(seq_length), device=device
        )
        z = mean + noise * (0.5 * logvar).exp()
        samples = self.vae.decode(z, vector, dense)
        if paste_observed and obs is not None:
            m = torch.ones_like(obs) if obs_mask is None else obs_mask.to(obs)
            samples = obs.to(samples) * m.to(samples) + samples * (1 - m.to(samples))
        return samples

    def configure_optimizers(self) -> Any:
        optimizer = torch.optim.Adam(
            self.parameters(),
            lr=self.config.lr,
            betas=(self.config.b1, self.config.b2),
        )
        return optimizer

reconstruct

reconstruct(x: Tensor) -> Tensor

Reconstruct input ECG signals (unconditional models).

Source code in src/ecgen/families/vae/vae_lightning.py

@torch.no_grad()
def reconstruct(self, x: Tensor) -> Tensor:
    """Reconstruct input ECG signals (unconditional models)."""
    self.vae.eval()
    recon, _, _ = self.vae(x)
    return recon

sample

sample(n_samples: int = 16, seq_length: int = 5000, cond: Optional[Mapping[str, Any]] = None, obs: Optional[Tensor] = None, obs_mask: Optional[Tensor] = None, paste_observed: bool = True, **kwargs: Any) -> Tensor

Sample ECGs: z from the prior (p(z | c) if learned), then decode.

Parameters:

Name	Type	Description	Default
n_samples	int	Number of ECGs.	16
seq_length	int	Length (multiple of the downsampling factor).	5000
cond	Optional[Mapping[str, Any]]	Raw condition values (see :meth:prepare_condition).	None
obs	Optional[Tensor]	Partially observed ECG [n, C, T].	None
obs_mask	Optional[Tensor]	Mask of obs (1 = observed).	None
paste_observed	bool	Copy the observed part of obs into the output.	True

Source code in src/ecgen/families/vae/vae_lightning.py

@torch.no_grad()
def sample(
    self,
    n_samples: int = 16,
    seq_length: int = 5000,
    cond: Optional[Mapping[str, Any]] = None,
    obs: Optional[Tensor] = None,
    obs_mask: Optional[Tensor] = None,
    paste_observed: bool = True,
    **kwargs: Any,
) -> Tensor:
    """
    Sample ECGs: ``z`` from the prior (``p(z | c)`` if learned), then decode.

    Args:
        n_samples: Number of ECGs.
        seq_length: Length (multiple of the downsampling factor).
        cond: Raw condition values (see :meth:`prepare_condition`).
        obs: Partially observed ECG ``[n, C, T]``.
        obs_mask: Mask of ``obs`` (1 = observed).
        paste_observed: Copy the observed part of ``obs`` into the output.
    """
    self.vae.eval()
    device = next(self.vae.parameters()).device
    if self.config.latent_type == "global":
        seq_length = self.config.seq_length
    latent_length = seq_length // self.vae.downsample_factor
    c: Optional[ConditionOutput] = self.encode_condition(
        n_samples, cond, obs, obs_mask
    )
    vector = c.vector if c is not None else None
    dense = c.dense if c is not None else None

    mean, logvar = self.vae.prior_params(
        latent_length, vector, dense, n_samples, device
    )
    noise = torch.randn(
        n_samples, *self.vae.latent_shape(seq_length), device=device
    )
    z = mean + noise * (0.5 * logvar).exp()
    samples = self.vae.decode(z, vector, dense)
    if paste_observed and obs is not None:
        m = torch.ones_like(obs) if obs_mask is None else obs_mask.to(obs)
        samples = obs.to(samples) * m.to(samples) + samples * (1 - m.to(samples))
    return samples

ecgen.families.vqvae.config

VQVAEConfig dataclass

Configuration for VQ-VAE model.

Source code in src/ecgen/families/vqvae/config.py

@dataclass
class VQVAEConfig:
    """Configuration for VQ-VAE model."""

    # Model architecture
    in_channels: int = 12
    base_channels: int = 64
    latent_channels: int = 64
    channel_multipliers: tuple[int, ...] = (1, 2, 4, 4)
    num_res_blocks: int = 2

    # Vector quantization
    num_embeddings: int = 512
    commitment_cost: float = 0.25

    # Training hyperparameters
    lr: float = 1e-4
    b1: float = 0.9
    b2: float = 0.999

PriorConfig dataclass

Configuration for PixelCNN prior model.

Source code in src/ecgen/families/vqvae/config.py

@dataclass
class PriorConfig:
    """Configuration for PixelCNN prior model."""

    # Model architecture
    num_embeddings: int = 512
    hidden_dim: int = 128
    num_layers: int = 10
    kernel_size: int = 2
    dilation_cycle: int = (
        10  # dilations 1, 2, 4, ..., 512 -> receptive field of 1024 codes
    )

    # Training hyperparameters
    lr: float = 1e-3
    b1: float = 0.9
    b2: float = 0.999

    # VQ-VAE checkpoint
    vqvae_checkpoint: str = ""
    # VQ-VAE architecture, filled in when vqvae_checkpoint is loaded so that
    # prior checkpoints can be restored without the stage-1 file
    vqvae_config: Optional[dict[str, Any]] = None

    # Conditional prior (the stage-1 VQ-VAE stays unconditional); see ecgen.conditioning
    conditioning: Optional[dict[str, Any]] = None

ecgen.families.vqvae.vector_quantizer

VectorQuantizer

Bases: Module

Vector Quantizer module for VQ-VAE.

Implements discrete latent space quantization using a learned codebook. Uses straight-through estimator for gradient flow during backpropagation.

Parameters:

Name	Type	Description	Default
num_embeddings	int	Size of the codebook (number of discrete codes)	512
embedding_dim	int	Dimensionality of each code vector	64
commitment_cost	float	Weight for commitment loss (encourages encoder outputs to commit to codes)	0.25

Source code in src/ecgen/families/vqvae/vector_quantizer.py

class VectorQuantizer(nn.Module):
    """
    Vector Quantizer module for VQ-VAE.

    Implements discrete latent space quantization using a learned codebook.
    Uses straight-through estimator for gradient flow during backpropagation.

    Args:
        num_embeddings: Size of the codebook (number of discrete codes)
        embedding_dim: Dimensionality of each code vector
        commitment_cost: Weight for commitment loss (encourages encoder outputs to commit to codes)
    """

    def __init__(
        self,
        num_embeddings: int = 512,
        embedding_dim: int = 64,
        commitment_cost: float = 0.25,
    ) -> None:
        super().__init__()
        self.num_embeddings = num_embeddings
        self.embedding_dim = embedding_dim
        self.commitment_cost = commitment_cost

        self.embedding = nn.Embedding(num_embeddings, embedding_dim)
        self.embedding.weight.data.uniform_(-1 / num_embeddings, 1 / num_embeddings)

    def forward(self, z: Tensor) -> tuple[Tensor, Tensor, Tensor]:
        """
        Quantize continuous latent vectors to discrete codes.

        Args:
            z: Encoder output [B, C, L] where C = embedding_dim

        Returns:
            vq_loss: Vector quantization loss (codebook + commitment)
            quantized: Quantized latent codes [B, C, L]
            indices: Codebook indices [B, L]
        """
        # Reshape: [B, C, L] -> [B, L, C]
        z = z.permute(0, 2, 1).contiguous()
        z_flattened = z.view(-1, self.embedding_dim)

        # Calculate L2 distances to all codebook vectors
        # ||z - e||^2 = ||z||^2 + ||e||^2 - 2 * z^T * e
        distances = (
            torch.sum(z_flattened**2, dim=1, keepdim=True)
            + torch.sum(self.embedding.weight**2, dim=1)
            - 2 * torch.matmul(z_flattened, self.embedding.weight.t())
        )

        # Find nearest codebook entry for each latent vector
        encoding_indices = torch.argmin(distances, dim=1).unsqueeze(1)
        encodings = torch.zeros(
            encoding_indices.shape[0], self.num_embeddings, device=z.device
        )
        encodings.scatter_(1, encoding_indices, 1)

        # Quantize by looking up codebook entries
        quantized = torch.matmul(encodings, self.embedding.weight).view(z.shape)

        # Loss: Codebook loss + Commitment loss
        # Codebook loss: move codebook vectors towards encoder outputs
        # Commitment loss: encourage encoder outputs to commit to chosen codes
        codebook_loss = F.mse_loss(quantized.detach(), z)
        commitment_loss = F.mse_loss(quantized, z.detach())
        vq_loss = codebook_loss + self.commitment_cost * commitment_loss

        # Straight-through estimator: copy gradients from quantized to z
        quantized = z + (quantized - z).detach()

        # Reshape back: [B, L, C] -> [B, C, L]
        quantized = quantized.permute(0, 2, 1).contiguous()

        # Return indices as [B, L]
        indices = encoding_indices.view(z.shape[0], -1)

        return vq_loss, quantized, indices

    @torch.no_grad()
    def get_codebook_entry(self, indices: Tensor) -> Tensor:
        """
        Get quantized vectors from codebook indices.

        Args:
            indices: Codebook indices [B, L]

        Returns:
            Quantized vectors [B, C, L]
        """
        # Get embeddings
        quantized = self.embedding(indices)  # [B, L, C]

        # Reshape to [B, C, L]
        quantized = quantized.permute(0, 2, 1).contiguous()

        return quantized

forward

forward(z: Tensor) -> tuple[Tensor, Tensor, Tensor]

Quantize continuous latent vectors to discrete codes.

Parameters:

Name	Type	Description	Default
z	Tensor	Encoder output [B, C, L] where C = embedding_dim	required

Returns:

Name	Type	Description
vq_loss	Tensor	Vector quantization loss (codebook + commitment)
quantized	Tensor	Quantized latent codes [B, C, L]
indices	Tensor	Codebook indices [B, L]

Source code in src/ecgen/families/vqvae/vector_quantizer.py

def forward(self, z: Tensor) -> tuple[Tensor, Tensor, Tensor]:
    """
    Quantize continuous latent vectors to discrete codes.

    Args:
        z: Encoder output [B, C, L] where C = embedding_dim

    Returns:
        vq_loss: Vector quantization loss (codebook + commitment)
        quantized: Quantized latent codes [B, C, L]
        indices: Codebook indices [B, L]
    """
    # Reshape: [B, C, L] -> [B, L, C]
    z = z.permute(0, 2, 1).contiguous()
    z_flattened = z.view(-1, self.embedding_dim)

    # Calculate L2 distances to all codebook vectors
    # ||z - e||^2 = ||z||^2 + ||e||^2 - 2 * z^T * e
    distances = (
        torch.sum(z_flattened**2, dim=1, keepdim=True)
        + torch.sum(self.embedding.weight**2, dim=1)
        - 2 * torch.matmul(z_flattened, self.embedding.weight.t())
    )

    # Find nearest codebook entry for each latent vector
    encoding_indices = torch.argmin(distances, dim=1).unsqueeze(1)
    encodings = torch.zeros(
        encoding_indices.shape[0], self.num_embeddings, device=z.device
    )
    encodings.scatter_(1, encoding_indices, 1)

    # Quantize by looking up codebook entries
    quantized = torch.matmul(encodings, self.embedding.weight).view(z.shape)

    # Loss: Codebook loss + Commitment loss
    # Codebook loss: move codebook vectors towards encoder outputs
    # Commitment loss: encourage encoder outputs to commit to chosen codes
    codebook_loss = F.mse_loss(quantized.detach(), z)
    commitment_loss = F.mse_loss(quantized, z.detach())
    vq_loss = codebook_loss + self.commitment_cost * commitment_loss

    # Straight-through estimator: copy gradients from quantized to z
    quantized = z + (quantized - z).detach()

    # Reshape back: [B, L, C] -> [B, C, L]
    quantized = quantized.permute(0, 2, 1).contiguous()

    # Return indices as [B, L]
    indices = encoding_indices.view(z.shape[0], -1)

    return vq_loss, quantized, indices

get_codebook_entry

get_codebook_entry(indices: Tensor) -> Tensor

Get quantized vectors from codebook indices.

Parameters:

Name	Type	Description	Default
indices	Tensor	Codebook indices [B, L]	required

Returns:

Type	Description
Tensor	Quantized vectors [B, C, L]

Source code in src/ecgen/families/vqvae/vector_quantizer.py

@torch.no_grad()
def get_codebook_entry(self, indices: Tensor) -> Tensor:
    """
    Get quantized vectors from codebook indices.

    Args:
        indices: Codebook indices [B, L]

    Returns:
        Quantized vectors [B, C, L]
    """
    # Get embeddings
    quantized = self.embedding(indices)  # [B, L, C]

    # Reshape to [B, C, L]
    quantized = quantized.permute(0, 2, 1).contiguous()

    return quantized

ecgen.families.vqvae.vqvae_1d

VQVAE1D

Bases: Module

Vector Quantized Variational Autoencoder for ECG signals.

Uses discrete latent representations instead of continuous ones. Architecture follows the existing VAE1D design but with vector quantization.

Parameters:

Name	Type	Description	Default
in_channels	int	Number of input channels (ECG leads)	12
base_channels	int	Base number of channels in encoder/decoder	64
latent_channels	int	Number of channels in latent space (embedding_dim)	64
channel_multipliers	tuple[int, ...]	Channel multipliers for each downsampling stage	(1, 2, 4, 4)
num_res_blocks	int	Number of residual blocks per stage	2
num_embeddings	int	Size of the discrete codebook	512
commitment_cost	float	Weight for commitment loss in VQ	0.25

Source code in src/ecgen/families/vqvae/vqvae_1d.py

class VQVAE1D(nn.Module):
    """
    Vector Quantized Variational Autoencoder for ECG signals.

    Uses discrete latent representations instead of continuous ones.
    Architecture follows the existing VAE1D design but with vector quantization.

    Args:
        in_channels: Number of input channels (ECG leads)
        base_channels: Base number of channels in encoder/decoder
        latent_channels: Number of channels in latent space (embedding_dim)
        channel_multipliers: Channel multipliers for each downsampling stage
        num_res_blocks: Number of residual blocks per stage
        num_embeddings: Size of the discrete codebook
        commitment_cost: Weight for commitment loss in VQ
    """

    def __init__(
        self,
        in_channels: int = 12,
        base_channels: int = 64,
        latent_channels: int = 64,
        channel_multipliers: tuple[int, ...] = (1, 2, 4, 4),
        num_res_blocks: int = 2,
        num_embeddings: int = 512,
        commitment_cost: float = 0.25,
    ) -> None:
        super().__init__()

        self.in_channels = in_channels
        self.latent_channels = latent_channels
        self.num_embeddings = num_embeddings

        # Encoder: ECG -> continuous latent
        self.encoder = Encoder1D(
            in_channels=in_channels,
            base_channels=base_channels,
            latent_channels=latent_channels,
            channel_multipliers=channel_multipliers,
            num_res_blocks=num_res_blocks,
        )

        # Vector Quantizer: continuous -> discrete
        self.vq = VectorQuantizer(
            num_embeddings=num_embeddings,
            embedding_dim=latent_channels,
            commitment_cost=commitment_cost,
        )

        # Decoder: discrete latent -> ECG
        self.decoder = Decoder1D(
            out_channels=in_channels,
            base_channels=base_channels,
            latent_channels=latent_channels,
            channel_multipliers=channel_multipliers,
            num_res_blocks=num_res_blocks,
        )

    def encode(self, x: Tensor) -> tuple[Tensor, Tensor]:
        """
        Encode input to continuous latent representation.

        Note: For VQ-VAE, encoder outputs a single tensor (not mean/logvar like VAE).
        We return it twice for API compatibility with VAE.

        Args:
            x: Input ECG signals [B, C, L]

        Returns:
            z: Continuous latent [B, latent_channels, L//downsample_factor]
            z: Same tensor (for compatibility)
        """
        z, _ = self.encoder(x)  # Encoder returns (mean, logvar), we only use mean
        return z, z

    def quantize(self, z: Tensor) -> tuple[Tensor, Tensor, Tensor]:
        """
        Quantize continuous latent to discrete codes.

        Args:
            z: Continuous latent [B, latent_channels, L_latent]

        Returns:
            vq_loss: Vector quantization loss
            z_q: Quantized latent [B, latent_channels, L_latent]
            indices: Codebook indices [B, L_latent]
        """
        return self.vq(z)

    def decode(self, z_q: Tensor) -> Tensor:
        """
        Decode quantized latent to reconstruction.

        Args:
            z_q: Quantized latent [B, latent_channels, L_latent]

        Returns:
            Reconstructed ECG [B, in_channels, L]
        """
        return self.decoder(z_q)

    def forward(self, x: Tensor) -> tuple[Tensor, Tensor, Tensor]:
        """
        Full forward pass: encode -> quantize -> decode.

        Args:
            x: Input ECG signals [B, in_channels, L]

        Returns:
            recon: Reconstructed ECG [B, in_channels, L]
            vq_loss: Vector quantization loss
            indices: Codebook indices [B, L_latent]
        """
        # Encode
        z, _ = self.encode(x)

        # Quantize
        vq_loss, z_q, indices = self.quantize(z)

        # Decode
        recon = self.decode(z_q)

        return recon, vq_loss, indices

    @torch.no_grad()
    def encode_to_indices(self, x: Tensor) -> Tensor:
        """
        Encode input ECG to discrete codebook indices.
        Used for training the prior model.

        Args:
            x: Input ECG signals [B, in_channels, L]

        Returns:
            indices: Codebook indices [B, L_latent]
        """
        z, _ = self.encode(x)
        _, _, indices = self.quantize(z)
        return indices

    @torch.no_grad()
    def decode_from_indices(self, indices: Tensor) -> Tensor:
        """
        Decode from discrete codebook indices to ECG.
        Used for generation from the prior model.

        Args:
            indices: Codebook indices [B, L_latent]

        Returns:
            Reconstructed ECG [B, in_channels, L]
        """
        z_q = self.vq.get_codebook_entry(indices)
        recon = self.decode(z_q)
        return recon

encode

encode(x: Tensor) -> tuple[Tensor, Tensor]

Encode input to continuous latent representation.

Note: For VQ-VAE, encoder outputs a single tensor (not mean/logvar like VAE). We return it twice for API compatibility with VAE.

Parameters:

Name	Type	Description	Default
x	Tensor	Input ECG signals [B, C, L]	required

Returns:

Name	Type	Description
z	Tensor	Continuous latent [B, latent_channels, L//downsample_factor]
z	Tensor	Same tensor (for compatibility)

Source code in src/ecgen/families/vqvae/vqvae_1d.py

def encode(self, x: Tensor) -> tuple[Tensor, Tensor]:
    """
    Encode input to continuous latent representation.

    Note: For VQ-VAE, encoder outputs a single tensor (not mean/logvar like VAE).
    We return it twice for API compatibility with VAE.

    Args:
        x: Input ECG signals [B, C, L]

    Returns:
        z: Continuous latent [B, latent_channels, L//downsample_factor]
        z: Same tensor (for compatibility)
    """
    z, _ = self.encoder(x)  # Encoder returns (mean, logvar), we only use mean
    return z, z

quantize

quantize(z: Tensor) -> tuple[Tensor, Tensor, Tensor]

Quantize continuous latent to discrete codes.

Parameters:

Name	Type	Description	Default
z	Tensor	Continuous latent [B, latent_channels, L_latent]	required

Returns:

Name	Type	Description
vq_loss	Tensor	Vector quantization loss
z_q	Tensor	Quantized latent [B, latent_channels, L_latent]
indices	Tensor	Codebook indices [B, L_latent]

Source code in src/ecgen/families/vqvae/vqvae_1d.py

def quantize(self, z: Tensor) -> tuple[Tensor, Tensor, Tensor]:
    """
    Quantize continuous latent to discrete codes.

    Args:
        z: Continuous latent [B, latent_channels, L_latent]

    Returns:
        vq_loss: Vector quantization loss
        z_q: Quantized latent [B, latent_channels, L_latent]
        indices: Codebook indices [B, L_latent]
    """
    return self.vq(z)

decode

decode(z_q: Tensor) -> Tensor

Decode quantized latent to reconstruction.

Parameters:

Name	Type	Description	Default
z_q	Tensor	Quantized latent [B, latent_channels, L_latent]	required

Returns:

Type	Description
Tensor	Reconstructed ECG [B, in_channels, L]

Source code in src/ecgen/families/vqvae/vqvae_1d.py

def decode(self, z_q: Tensor) -> Tensor:
    """
    Decode quantized latent to reconstruction.

    Args:
        z_q: Quantized latent [B, latent_channels, L_latent]

    Returns:
        Reconstructed ECG [B, in_channels, L]
    """
    return self.decoder(z_q)

forward

forward(x: Tensor) -> tuple[Tensor, Tensor, Tensor]

Full forward pass: encode -> quantize -> decode.

Parameters:

Name	Type	Description	Default
x	Tensor	Input ECG signals [B, in_channels, L]	required

Returns:

Name	Type	Description
recon	Tensor	Reconstructed ECG [B, in_channels, L]
vq_loss	Tensor	Vector quantization loss
indices	Tensor	Codebook indices [B, L_latent]

Source code in src/ecgen/families/vqvae/vqvae_1d.py

def forward(self, x: Tensor) -> tuple[Tensor, Tensor, Tensor]:
    """
    Full forward pass: encode -> quantize -> decode.

    Args:
        x: Input ECG signals [B, in_channels, L]

    Returns:
        recon: Reconstructed ECG [B, in_channels, L]
        vq_loss: Vector quantization loss
        indices: Codebook indices [B, L_latent]
    """
    # Encode
    z, _ = self.encode(x)

    # Quantize
    vq_loss, z_q, indices = self.quantize(z)

    # Decode
    recon = self.decode(z_q)

    return recon, vq_loss, indices

encode_to_indices

encode_to_indices(x: Tensor) -> Tensor

Encode input ECG to discrete codebook indices. Used for training the prior model.

Parameters:

Name	Type	Description	Default
x	Tensor	Input ECG signals [B, in_channels, L]	required

Returns:

Name	Type	Description
indices	Tensor	Codebook indices [B, L_latent]

Source code in src/ecgen/families/vqvae/vqvae_1d.py

@torch.no_grad()
def encode_to_indices(self, x: Tensor) -> Tensor:
    """
    Encode input ECG to discrete codebook indices.
    Used for training the prior model.

    Args:
        x: Input ECG signals [B, in_channels, L]

    Returns:
        indices: Codebook indices [B, L_latent]
    """
    z, _ = self.encode(x)
    _, _, indices = self.quantize(z)
    return indices

decode_from_indices

decode_from_indices(indices: Tensor) -> Tensor

Decode from discrete codebook indices to ECG. Used for generation from the prior model.

Parameters:

Name	Type	Description	Default
indices	Tensor	Codebook indices [B, L_latent]	required

Returns:

Type	Description
Tensor	Reconstructed ECG [B, in_channels, L]

Source code in src/ecgen/families/vqvae/vqvae_1d.py

@torch.no_grad()
def decode_from_indices(self, indices: Tensor) -> Tensor:
    """
    Decode from discrete codebook indices to ECG.
    Used for generation from the prior model.

    Args:
        indices: Codebook indices [B, L_latent]

    Returns:
        Reconstructed ECG [B, in_channels, L]
    """
    z_q = self.vq.get_codebook_entry(indices)
    recon = self.decode(z_q)
    return recon

ecgen.families.vqvae.vqvae_lightning

VQVAELightning

Bases: GenerativeModule

PyTorch Lightning wrapper for VQ-VAE model.

Handles training, validation, and sampling for Stage 1 (VQ-VAE training).

Source code in src/ecgen/families/vqvae/vqvae_lightning.py

@FAMILIES.register("vqvae")
class VQVAELightning(GenerativeModule):
    """
    PyTorch Lightning wrapper for VQ-VAE model.

    Handles training, validation, and sampling for Stage 1 (VQ-VAE training).
    """

    family = "vqvae"
    supports = Capabilities(stages=2)

    def __init__(
        self, config: VQVAEConfig | dict[str, Any] | None = None, **kwargs: Any
    ) -> None:
        super().__init__()
        config = resolve_config(VQVAEConfig, config, kwargs)
        self.save_hyperparameters(config.__dict__)
        self.config = config

        self.vqvae = VQVAE1D(
            in_channels=config.in_channels,
            base_channels=config.base_channels,
            latent_channels=config.latent_channels,
            channel_multipliers=config.channel_multipliers,
            num_res_blocks=config.num_res_blocks,
            num_embeddings=config.num_embeddings,
            commitment_cost=config.commitment_cost,
        )

        self._val_real_sample: Tensor | None = None
        self._val_recon_sample: Tensor | None = None

    def forward(self, x: Tensor) -> tuple[Tensor, Tensor, Tensor]:
        return self.vqvae(x)

    def training_step(self, batch: Any, batch_idx: int) -> Tensor:
        ecgs = self.get_signal(batch)

        recon, vq_loss_val, indices = self.vqvae(ecgs)
        total_loss, recon_loss, vq_loss_component = vqvae_loss(recon, ecgs, vq_loss_val)

        self.log(
            "train/total_loss", total_loss, on_step=True, on_epoch=True, prog_bar=True
        )
        self.log(
            "train/recon_loss", recon_loss, on_step=True, on_epoch=True, prog_bar=True
        )
        self.log(
            "train/vq_loss",
            vq_loss_component,
            on_step=True,
            on_epoch=True,
            prog_bar=False,
        )

        # Log codebook usage statistics
        if batch_idx % 100 == 0:
            unique_codes = torch.unique(indices).numel()
            self.log(
                "train/unique_codes", float(unique_codes), on_step=True, on_epoch=False
            )
            self.log(
                "train/codebook_usage",
                float(unique_codes) / self.config.num_embeddings,
                on_step=True,
                on_epoch=False,
            )

        return total_loss

    def validation_step(self, batch: Any, batch_idx: int) -> Tensor:
        ecgs = self.get_signal(batch)

        recon, vq_loss_val, indices = self.vqvae(ecgs)
        total_loss, recon_loss, vq_loss_component = vqvae_loss(recon, ecgs, vq_loss_val)

        self.log(
            "val/total_loss", total_loss, on_step=False, on_epoch=True, prog_bar=True
        )
        self.log(
            "val/recon_loss", recon_loss, on_step=False, on_epoch=True, prog_bar=True
        )
        self.log(
            "val/vq_loss",
            vq_loss_component,
            on_step=False,
            on_epoch=True,
            prog_bar=False,
        )
        self.log("val_loss", total_loss, on_step=False, on_epoch=True, prog_bar=True)

        # Log codebook usage
        unique_codes = torch.unique(indices).numel()
        self.log("val/unique_codes", float(unique_codes), on_step=False, on_epoch=True)
        self.log(
            "val/codebook_usage",
            float(unique_codes) / self.config.num_embeddings,
            on_step=False,
            on_epoch=True,
        )

        if batch_idx == 0:
            self._val_real_sample = ecgs[0].detach().cpu()
            self._val_recon_sample = recon[0].detach().cpu()

        return total_loss

    @property
    def downsample_factor(self) -> int:
        """Temporal downsampling between ECG and latent (one stride-2 stage per multiplier after the first)."""
        return 2 ** (len(self.config.channel_multipliers) - 1)

    def latent_length(self, seq_length: int) -> int:
        """Number of latent codes for an ECG of ``seq_length`` samples."""
        if seq_length % self.downsample_factor != 0:
            raise ValueError(
                f"seq_length={seq_length} must be a multiple of the downsampling factor "
                f"{self.downsample_factor}; decoded samples would be shorter"
            )
        return seq_length // self.downsample_factor

    @torch.no_grad()
    def reconstruct(self, x: Tensor) -> Tensor:
        """Reconstruct input ECG signals."""
        self.vqvae.eval()
        recon, _, _ = self.vqvae(x)
        return recon

    @torch.no_grad()
    def encode_to_indices(self, x: Tensor) -> Tensor:
        """
        Encode ECG to discrete codebook indices.
        Used for training the prior model.
        """
        self.vqvae.eval()
        return self.vqvae.encode_to_indices(x)

    @torch.no_grad()
    def decode_from_indices(self, indices: Tensor) -> Tensor:
        """
        Decode from discrete codebook indices.
        Used for generation from the prior model.
        """
        self.vqvae.eval()
        return self.vqvae.decode_from_indices(indices)

    @torch.no_grad()
    def sample(self, n_samples: int = 16, seq_length: int = 5000) -> Tensor:
        """
        Sample by randomly selecting codebook indices.

        Note: This is not a good sampling method. For proper generation,
        use the trained prior model (Stage 2).
        """
        self.vqvae.eval()
        device = next(self.vqvae.parameters()).device

        latent_length = self.latent_length(seq_length)

        # Random indices (not a good prior!)
        indices = torch.randint(
            0,
            self.config.num_embeddings,
            (n_samples, latent_length),
            device=device,
        )

        samples = self.vqvae.decode_from_indices(indices)
        return samples

    def configure_optimizers(self) -> Any:
        optimizer = torch.optim.Adam(
            self.vqvae.parameters(),
            lr=self.config.lr,
            betas=(self.config.b1, self.config.b2),
        )
        return optimizer

downsample_factor property

downsample_factor: int

Temporal downsampling between ECG and latent (one stride-2 stage per multiplier after the first).

latent_length

latent_length(seq_length: int) -> int

Number of latent codes for an ECG of seq_length samples.

Source code in src/ecgen/families/vqvae/vqvae_lightning.py

def latent_length(self, seq_length: int) -> int:
    """Number of latent codes for an ECG of ``seq_length`` samples."""
    if seq_length % self.downsample_factor != 0:
        raise ValueError(
            f"seq_length={seq_length} must be a multiple of the downsampling factor "
            f"{self.downsample_factor}; decoded samples would be shorter"
        )
    return seq_length // self.downsample_factor

reconstruct

reconstruct(x: Tensor) -> Tensor

Reconstruct input ECG signals.

Source code in src/ecgen/families/vqvae/vqvae_lightning.py

@torch.no_grad()
def reconstruct(self, x: Tensor) -> Tensor:
    """Reconstruct input ECG signals."""
    self.vqvae.eval()
    recon, _, _ = self.vqvae(x)
    return recon

encode_to_indices

encode_to_indices(x: Tensor) -> Tensor

Encode ECG to discrete codebook indices. Used for training the prior model.

Source code in src/ecgen/families/vqvae/vqvae_lightning.py

@torch.no_grad()
def encode_to_indices(self, x: Tensor) -> Tensor:
    """
    Encode ECG to discrete codebook indices.
    Used for training the prior model.
    """
    self.vqvae.eval()
    return self.vqvae.encode_to_indices(x)

decode_from_indices

decode_from_indices(indices: Tensor) -> Tensor

Decode from discrete codebook indices. Used for generation from the prior model.

Source code in src/ecgen/families/vqvae/vqvae_lightning.py

@torch.no_grad()
def decode_from_indices(self, indices: Tensor) -> Tensor:
    """
    Decode from discrete codebook indices.
    Used for generation from the prior model.
    """
    self.vqvae.eval()
    return self.vqvae.decode_from_indices(indices)

sample

sample(n_samples: int = 16, seq_length: int = 5000) -> Tensor

Sample by randomly selecting codebook indices.

Note: This is not a good sampling method. For proper generation, use the trained prior model (Stage 2).

Source code in src/ecgen/families/vqvae/vqvae_lightning.py

@torch.no_grad()
def sample(self, n_samples: int = 16, seq_length: int = 5000) -> Tensor:
    """
    Sample by randomly selecting codebook indices.

    Note: This is not a good sampling method. For proper generation,
    use the trained prior model (Stage 2).
    """
    self.vqvae.eval()
    device = next(self.vqvae.parameters()).device

    latent_length = self.latent_length(seq_length)

    # Random indices (not a good prior!)
    indices = torch.randint(
        0,
        self.config.num_embeddings,
        (n_samples, latent_length),
        device=device,
    )

    samples = self.vqvae.decode_from_indices(indices)
    return samples

ecgen.families.vqvae.pixelcnn_prior

CausalGatedBlock

Bases: Module

Residual gated causal 1D convolution (WaveNet-style).

Output at position t depends only on inputs at positions <= t: the input is left-padded by (kernel_size - 1) * dilation, so no future positions are seen. Uses gated activation: tanh(W_f * x) * sigmoid(W_g * x)

Parameters:

Name	Type	Description	Default
channels	int	Number of input/output channels	required
kernel_size	int	Size of convolutional kernel	2
dilation	int	Dilation of the convolution	1

Source code in src/ecgen/families/vqvae/pixelcnn_prior.py

class CausalGatedBlock(nn.Module):
    """
    Residual gated causal 1D convolution (WaveNet-style).

    Output at position t depends only on inputs at positions <= t: the input is
    left-padded by (kernel_size - 1) * dilation, so no future positions are seen.
    Uses gated activation: tanh(W_f * x) * sigmoid(W_g * x)

    Args:
        channels: Number of input/output channels
        kernel_size: Size of convolutional kernel
        dilation: Dilation of the convolution
    """

    def __init__(
        self,
        channels: int,
        kernel_size: int = 2,
        dilation: int = 1,
        cond_dim: int = 0,
        dense_channels: int = 0,
    ) -> None:
        super().__init__()
        self.left_pad = (kernel_size - 1) * dilation
        # Conditional PixelCNN (van den Oord et al., 2016): the condition adds
        # V_f h / V_g h to both gate pre-activations; zero-initialised
        self.cond_proj = nn.Linear(cond_dim, 2 * channels) if cond_dim else None
        self.dense_proj = (
            nn.Conv1d(dense_channels, 2 * channels, 1) if dense_channels else None
        )
        for m in (self.cond_proj, self.dense_proj):
            if m is not None:
                nn.init.zeros_(m.weight)
                nn.init.zeros_(m.bias)

        # Gated convolution: outputs 2x channels for tanh and sigmoid gates
        self.conv = nn.Conv1d(channels, 2 * channels, kernel_size, dilation=dilation)
        self.residual = nn.Conv1d(channels, channels, kernel_size=1)
        self.skip = nn.Conv1d(channels, channels, kernel_size=1)

    def forward(
        self, x: Tensor, c: Optional[Tensor] = None, dense: Optional[Tensor] = None
    ) -> tuple[Tensor, Tensor]:
        """
        Args:
            x: Input tensor [B, C, L]

        Returns:
            (residual output [B, C, L], skip output [B, C, L])
        """
        out = self.conv(F.pad(x, (self.left_pad, 0)))
        if self.cond_proj is not None and c is not None:
            out = out + self.cond_proj(c)[..., None]
        if self.dense_proj is not None and dense is not None:
            out = out + self.dense_proj(dense)

        # Split into two halves for gated activation
        out_tanh, out_sigmoid = torch.chunk(out, 2, dim=1)
        h = torch.tanh(out_tanh) * torch.sigmoid(out_sigmoid)

        return x + self.residual(h), self.skip(h)

forward

forward(x: Tensor, c: Optional[Tensor] = None, dense: Optional[Tensor] = None) -> tuple[Tensor, Tensor]

Parameters:

Name	Type	Description	Default
x	Tensor	Input tensor [B, C, L]	required

Returns:

Type	Description
tuple[Tensor, Tensor]	(residual output [B, C, L], skip output [B, C, L])

Source code in src/ecgen/families/vqvae/pixelcnn_prior.py

def forward(
    self, x: Tensor, c: Optional[Tensor] = None, dense: Optional[Tensor] = None
) -> tuple[Tensor, Tensor]:
    """
    Args:
        x: Input tensor [B, C, L]

    Returns:
        (residual output [B, C, L], skip output [B, C, L])
    """
    out = self.conv(F.pad(x, (self.left_pad, 0)))
    if self.cond_proj is not None and c is not None:
        out = out + self.cond_proj(c)[..., None]
    if self.dense_proj is not None and dense is not None:
        out = out + self.dense_proj(dense)

    # Split into two halves for gated activation
    out_tanh, out_sigmoid = torch.chunk(out, 2, dim=1)
    h = torch.tanh(out_tanh) * torch.sigmoid(out_sigmoid)

    return x + self.residual(h), self.skip(h)

PixelCNNPrior

Bases: Module

Autoregressive prior for VQ-VAE discrete codes.

Learns the distribution p(z) over discrete latent codes. The embedded codes are shifted right by one position, then passed through causal dilated convolutions, so the logits at position t depend only on codes < t and predict code t. Dilations double every layer (1, 2, 4, ...) and restart every dilation_cycle layers, so the receptive field grows exponentially with depth (see receptive_field).

Parameters:

Name	Type	Description	Default
num_embeddings	int	Size of the codebook (vocabulary size)	512
hidden_dim	int	Hidden dimension for the network	128
num_layers	int	Number of gated causal conv layers	10
kernel_size	int	Kernel size of the causal convolutions	2
dilation_cycle	int	Number of layers after which the dilation resets to 1	10

Source code in src/ecgen/families/vqvae/pixelcnn_prior.py

class PixelCNNPrior(nn.Module):
    """
    Autoregressive prior for VQ-VAE discrete codes.

    Learns the distribution p(z) over discrete latent codes. The embedded codes
    are shifted right by one position, then passed through causal dilated
    convolutions, so the logits at position t depend only on codes < t and
    predict code t. Dilations double every layer (1, 2, 4, ...) and restart
    every ``dilation_cycle`` layers, so the receptive field grows exponentially
    with depth (see ``receptive_field``).

    Args:
        num_embeddings: Size of the codebook (vocabulary size)
        hidden_dim: Hidden dimension for the network
        num_layers: Number of gated causal conv layers
        kernel_size: Kernel size of the causal convolutions
        dilation_cycle: Number of layers after which the dilation resets to 1
    """

    def __init__(
        self,
        num_embeddings: int = 512,
        hidden_dim: int = 128,
        num_layers: int = 10,
        kernel_size: int = 2,
        dilation_cycle: int = 10,
        cond_dim: int = 0,
        dense_channels: int = 0,
    ) -> None:
        super().__init__()
        self.num_embeddings = num_embeddings
        self.hidden_dim = hidden_dim

        # Embedding layer: convert discrete indices to continuous vectors
        self.embedding = nn.Embedding(num_embeddings, hidden_dim)

        self.dilations = [2 ** (i % dilation_cycle) for i in range(num_layers)]
        self.kernel_size = kernel_size
        self.blocks = nn.ModuleList(
            CausalGatedBlock(hidden_dim, kernel_size, d, cond_dim, dense_channels)
            for d in self.dilations
        )

        # Output head: project summed skip connections to logits over codebook
        self.head = nn.Sequential(
            nn.ReLU(),
            nn.Conv1d(hidden_dim, hidden_dim, kernel_size=1),
            nn.ReLU(),
            nn.Conv1d(hidden_dim, num_embeddings, kernel_size=1),
        )

    @property
    def receptive_field(self) -> int:
        """Number of preceding codes that can influence the prediction of a code."""
        return 1 + sum((self.kernel_size - 1) * d for d in self.dilations)

    def forward(
        self,
        indices: Tensor,
        c: Optional[Tensor] = None,
        dense: Optional[Tensor] = None,
    ) -> Tensor:
        """
        Compute logits for each discrete code given the codes before it.

        Args:
            indices: Discrete code indices [B, L]
            c: Optional global condition [B, cond_dim]
            dense: Optional condition at code resolution [B, dense_channels, L]

        Returns:
            Logits over codebook [B, num_embeddings, L]; position t predicts
            indices[:, t] from indices[:, :t]
        """
        # Embed indices: [B, L] -> [B, hidden_dim, L]
        x = self.embedding(indices).permute(0, 2, 1)

        # Shift right by one so position t only sees codes < t (zeros at t = 0)
        x = F.pad(x, (1, 0))[:, :, :-1]

        skip = torch.zeros_like(x)
        for block in self.blocks:
            x, s = block(x, c, dense)
            skip = skip + s

        return self.head(skip)

    @torch.no_grad()
    def sample(
        self,
        batch_size: int = 1,
        latent_length: int = 625,  # 5000 samples / 8x VQ-VAE downsampling
        temperature: float = 1.0,
        device: str = "cuda",
        c: Optional[Tensor] = None,
        dense: Optional[Tensor] = None,
        c_uncond: Optional[Tensor] = None,
        dense_uncond: Optional[Tensor] = None,
        guidance: float = 1.0,
    ) -> Tensor:
        """
        Sample discrete codes autoregressively from the prior.

        Args:
            batch_size: Number of samples to generate
            latent_length: Length of latent sequence
            temperature: Sampling temperature (higher = more random)
            device: Device to generate samples on
            c: Global condition (see :meth:`forward`).
            dense: Condition at code resolution, covering all ``latent_length`` positions.
            c_uncond: Null global condition, for classifier-free guidance.
            dense_uncond: Null dense condition, for classifier-free guidance.
            guidance: CFG scale ``w``; logits = ``l_u + w (l_c - l_u)``. 1 = no guidance.

        Returns:
            Sampled indices [B, L]
        """
        indices = torch.zeros(
            batch_size, latent_length, dtype=torch.long, device=device
        )
        use_cfg = guidance != 1.0 and (c_uncond is not None or dense_uncond is not None)

        def window(d: Optional[Tensor], start: int, stop: int) -> Optional[Tensor]:
            return None if d is None else d[..., start:stop]

        for i in range(latent_length):
            # Logits at position i depend only on codes < i, so only the last
            # receptive_field codes before i are needed
            start = max(0, i - self.receptive_field)
            ctx = indices[:, start : i + 1]
            logits = self.forward(ctx, c, window(dense, start, i + 1))[:, :, -1]
            if use_cfg:
                uncond = self.forward(
                    ctx, c_uncond, window(dense_uncond, start, i + 1)
                )[:, :, -1]
                logits = uncond + guidance * (logits - uncond)
            logits = logits / temperature

            # Sample from categorical distribution
            probs = torch.softmax(logits, dim=-1)
            indices[:, i] = torch.multinomial(probs, num_samples=1).squeeze(-1)

        return indices

receptive_field property

receptive_field: int

Number of preceding codes that can influence the prediction of a code.

forward

forward(indices: Tensor, c: Optional[Tensor] = None, dense: Optional[Tensor] = None) -> Tensor

Compute logits for each discrete code given the codes before it.

Parameters:

Name	Type	Description	Default
indices	Tensor	Discrete code indices [B, L]	required
c	Optional[Tensor]	Optional global condition [B, cond_dim]	None
dense	Optional[Tensor]	Optional condition at code resolution [B, dense_channels, L]	None

Returns:

Type	Description
Tensor	Logits over codebook [B, num_embeddings, L]; position t predicts
Tensor	indices[:, t] from indices[:, :t]

Source code in src/ecgen/families/vqvae/pixelcnn_prior.py

def forward(
    self,
    indices: Tensor,
    c: Optional[Tensor] = None,
    dense: Optional[Tensor] = None,
) -> Tensor:
    """
    Compute logits for each discrete code given the codes before it.

    Args:
        indices: Discrete code indices [B, L]
        c: Optional global condition [B, cond_dim]
        dense: Optional condition at code resolution [B, dense_channels, L]

    Returns:
        Logits over codebook [B, num_embeddings, L]; position t predicts
        indices[:, t] from indices[:, :t]
    """
    # Embed indices: [B, L] -> [B, hidden_dim, L]
    x = self.embedding(indices).permute(0, 2, 1)

    # Shift right by one so position t only sees codes < t (zeros at t = 0)
    x = F.pad(x, (1, 0))[:, :, :-1]

    skip = torch.zeros_like(x)
    for block in self.blocks:
        x, s = block(x, c, dense)
        skip = skip + s

    return self.head(skip)

sample

sample(batch_size: int = 1, latent_length: int = 625, temperature: float = 1.0, device: str = 'cuda', c: Optional[Tensor] = None, dense: Optional[Tensor] = None, c_uncond: Optional[Tensor] = None, dense_uncond: Optional[Tensor] = None, guidance: float = 1.0) -> Tensor

Sample discrete codes autoregressively from the prior.

Parameters:

Name	Type	Description	Default
batch_size	int	Number of samples to generate	1
latent_length	int	Length of latent sequence	625
temperature	float	Sampling temperature (higher = more random)	1.0
device	str	Device to generate samples on	'cuda'
c	Optional[Tensor]	Global condition (see :meth:forward).	None
dense	Optional[Tensor]	Condition at code resolution, covering all latent_length positions.	None
c_uncond	Optional[Tensor]	Null global condition, for classifier-free guidance.	None
dense_uncond	Optional[Tensor]	Null dense condition, for classifier-free guidance.	None
guidance	float	CFG scale w; logits = l_u + w (l_c - l_u). 1 = no guidance.	1.0

Returns:

Type	Description
Tensor	Sampled indices [B, L]

Source code in src/ecgen/families/vqvae/pixelcnn_prior.py

@torch.no_grad()
def sample(
    self,
    batch_size: int = 1,
    latent_length: int = 625,  # 5000 samples / 8x VQ-VAE downsampling
    temperature: float = 1.0,
    device: str = "cuda",
    c: Optional[Tensor] = None,
    dense: Optional[Tensor] = None,
    c_uncond: Optional[Tensor] = None,
    dense_uncond: Optional[Tensor] = None,
    guidance: float = 1.0,
) -> Tensor:
    """
    Sample discrete codes autoregressively from the prior.

    Args:
        batch_size: Number of samples to generate
        latent_length: Length of latent sequence
        temperature: Sampling temperature (higher = more random)
        device: Device to generate samples on
        c: Global condition (see :meth:`forward`).
        dense: Condition at code resolution, covering all ``latent_length`` positions.
        c_uncond: Null global condition, for classifier-free guidance.
        dense_uncond: Null dense condition, for classifier-free guidance.
        guidance: CFG scale ``w``; logits = ``l_u + w (l_c - l_u)``. 1 = no guidance.

    Returns:
        Sampled indices [B, L]
    """
    indices = torch.zeros(
        batch_size, latent_length, dtype=torch.long, device=device
    )
    use_cfg = guidance != 1.0 and (c_uncond is not None or dense_uncond is not None)

    def window(d: Optional[Tensor], start: int, stop: int) -> Optional[Tensor]:
        return None if d is None else d[..., start:stop]

    for i in range(latent_length):
        # Logits at position i depend only on codes < i, so only the last
        # receptive_field codes before i are needed
        start = max(0, i - self.receptive_field)
        ctx = indices[:, start : i + 1]
        logits = self.forward(ctx, c, window(dense, start, i + 1))[:, :, -1]
        if use_cfg:
            uncond = self.forward(
                ctx, c_uncond, window(dense_uncond, start, i + 1)
            )[:, :, -1]
            logits = uncond + guidance * (logits - uncond)
        logits = logits / temperature

        # Sample from categorical distribution
        probs = torch.softmax(logits, dim=-1)
        indices[:, i] = torch.multinomial(probs, num_samples=1).squeeze(-1)

    return indices

ecgen.families.vqvae.prior_lightning

PriorLightning

Bases: GenerativeModule

PyTorch Lightning wrapper for PixelCNN prior model.

Handles training and sampling for Stage 2 (prior training). Requires a trained VQ-VAE checkpoint to extract discrete codes.

Source code in src/ecgen/families/vqvae/prior_lightning.py

@FAMILIES.register("vqvae_prior")
class PriorLightning(GenerativeModule):
    """
    PyTorch Lightning wrapper for PixelCNN prior model.

    Handles training and sampling for Stage 2 (prior training).
    Requires a trained VQ-VAE checkpoint to extract discrete codes.
    """

    family = "vqvae"
    supports = Capabilities(conditional=True, cfg=True, likelihood=True, stages=2)
    accepts = frozenset({"vector", "dense"})

    def __init__(
        self, config: PriorConfig | dict[str, Any] | None = None, **kwargs: Any
    ) -> None:
        super().__init__()
        config = resolve_config(PriorConfig, config, kwargs)
        self.save_hyperparameters(config.__dict__)
        self.config = config

        self.init_conditioning(config.conditioning)
        cond_dim = dense_channels = 0
        if self.conditioner is not None:
            cond_dim = (
                self.conditioner.vector_dim if "vector" in self.conditioner.slots else 0
            )
            dense_channels = self.conditioner.dense_channels

        self.prior = PixelCNNPrior(
            num_embeddings=config.num_embeddings,
            hidden_dim=config.hidden_dim,
            num_layers=config.num_layers,
            kernel_size=config.kernel_size,
            dilation_cycle=config.dilation_cycle,
            cond_dim=cond_dim,
            dense_channels=dense_channels,
        )

        # Frozen VQ-VAE for encoding/decoding. When restoring a prior checkpoint,
        # vqvae_config is set and the VQ-VAE weights come from this checkpoint's
        # state_dict, so the original stage-1 file is not needed.
        self.vqvae: VQVAELightning | None = None
        if config.vqvae_config is not None:
            self.vqvae = VQVAELightning(config.vqvae_config)
            self.vqvae.eval()
            self.vqvae.freeze()
        elif config.vqvae_checkpoint:
            self._load_vqvae(config.vqvae_checkpoint)

    def _load_vqvae(self, checkpoint_path: str) -> None:
        """Load frozen VQ-VAE model from checkpoint."""
        if not Path(checkpoint_path).exists():
            raise FileNotFoundError(f"VQ-VAE checkpoint not found: {checkpoint_path}")

        print(f"Loading VQ-VAE from: {checkpoint_path}")
        self.vqvae = VQVAELightning.load_from_checkpoint(
            checkpoint_path, map_location="cpu"
        )
        self.vqvae.eval()
        self.vqvae.freeze()
        # Store the architecture so prior checkpoints are self-contained
        vqvae_config = dict(self.vqvae.config.__dict__)
        self.config.vqvae_config = vqvae_config
        self.hparams["vqvae_config"] = vqvae_config
        print("VQ-VAE loaded and frozen")

    def forward(self, indices: Tensor) -> Tensor:
        return self.prior(indices)

    def compute_loss(
        self,
        indices: Tensor,
        c: Optional[Tensor] = None,
        dense: Optional[Tensor] = None,
    ) -> Tensor:
        """Cross-entropy of each code given the codes before it (logits at t predict code t)."""
        logits = self.prior(indices, c, dense)  # [B, num_embeddings, L]
        return prior_loss(logits, indices)

    @staticmethod
    def _latent_dense(dense: Optional[Tensor], length: int) -> Optional[Tensor]:
        return None if dense is None else F.adaptive_avg_pool1d(dense, length)

    def _batch_loss(self, batch: Any) -> Tensor:
        ecgs = self.get_signal(batch)
        # Extract discrete codes from frozen VQ-VAE
        with torch.no_grad():
            if self.vqvae is None:
                raise RuntimeError(
                    "VQ-VAE not loaded. Provide vqvae_checkpoint in config."
                )
            indices = self.vqvae.encode_to_indices(ecgs)  # [B, L]
        cond = self.condition(batch)
        if cond is None:
            return self.compute_loss(indices)
        return self.compute_loss(
            indices, cond.vector, self._latent_dense(cond.dense, indices.shape[1])
        )

    def training_step(self, batch: Any, batch_idx: int) -> Tensor:
        loss = self._batch_loss(batch)

        self.log("train/loss", loss, on_step=True, on_epoch=True, prog_bar=True)

        return loss

    def validation_step(self, batch: Any, batch_idx: int) -> Tensor:
        loss = self._batch_loss(batch)

        self.log("val/loss", loss, on_step=False, on_epoch=True, prog_bar=True)
        self.log("val_loss", loss, on_step=False, on_epoch=True, prog_bar=True)

        return loss

    @torch.no_grad()
    def sample(
        self,
        n_samples: int = 16,
        seq_length: int = 5000,
        temperature: float = 1.0,
        cond: Optional[Mapping[str, Any]] = None,
        obs: Optional[Tensor] = None,
        obs_mask: Optional[Tensor] = None,
        guidance: float = 1.0,
        **kwargs: Any,
    ) -> Tensor:
        """
        Generate ECG samples by sampling from the prior and decoding.

        Args:
            n_samples: Number of samples to generate
            seq_length: Length of ECG sequence
            temperature: Sampling temperature
            cond: Raw condition values (conditional priors)
            obs: Partially observed ECG (dense conditions)
            obs_mask: Mask of ``obs`` (1 = observed)
            guidance: Classifier-free guidance scale (1 = off)

        Returns:
            Generated ECG samples [B, C, L]
        """
        if self.vqvae is None:
            raise RuntimeError("VQ-VAE not loaded. Cannot decode samples.")

        self.prior.eval()
        device = next(self.prior.parameters()).device

        latent_length = self.vqvae.latent_length(seq_length)

        kwargs_c: dict[str, Any] = {}
        c = self.encode_condition(n_samples, cond, obs, obs_mask)
        if c is not None:
            kwargs_c = {
                "c": c.vector,
                "dense": self._latent_dense(c.dense, latent_length),
            }
            if guidance != 1.0:
                u = self.encode_condition(n_samples, cond, obs, obs_mask, drop="all")
                kwargs_c["c_uncond"] = u.vector
                kwargs_c["dense_uncond"] = self._latent_dense(u.dense, latent_length)
                kwargs_c["guidance"] = guidance

        # Sample discrete codes from prior
        indices = self.prior.sample(
            batch_size=n_samples,
            latent_length=latent_length,
            temperature=temperature,
            device=device,
            **kwargs_c,
        )

        # Decode to ECG
        samples = self.vqvae.decode_from_indices(indices)

        return samples

    def configure_optimizers(self) -> Any:
        params = list(self.prior.parameters())
        if self.conditioner is not None:
            params += list(self.conditioner.parameters())
        optimizer = torch.optim.Adam(
            params,
            lr=self.config.lr,
            betas=(self.config.b1, self.config.b2),
        )
        return optimizer

compute_loss

compute_loss(indices: Tensor, c: Optional[Tensor] = None, dense: Optional[Tensor] = None) -> Tensor

Cross-entropy of each code given the codes before it (logits at t predict code t).

Source code in src/ecgen/families/vqvae/prior_lightning.py

def compute_loss(
    self,
    indices: Tensor,
    c: Optional[Tensor] = None,
    dense: Optional[Tensor] = None,
) -> Tensor:
    """Cross-entropy of each code given the codes before it (logits at t predict code t)."""
    logits = self.prior(indices, c, dense)  # [B, num_embeddings, L]
    return prior_loss(logits, indices)

sample

sample(n_samples: int = 16, seq_length: int = 5000, temperature: float = 1.0, cond: Optional[Mapping[str, Any]] = None, obs: Optional[Tensor] = None, obs_mask: Optional[Tensor] = None, guidance: float = 1.0, **kwargs: Any) -> Tensor

Generate ECG samples by sampling from the prior and decoding.

Parameters:

Name	Type	Description	Default
n_samples	int	Number of samples to generate	16
seq_length	int	Length of ECG sequence	5000
temperature	float	Sampling temperature	1.0
cond	Optional[Mapping[str, Any]]	Raw condition values (conditional priors)	None
obs	Optional[Tensor]	Partially observed ECG (dense conditions)	None
obs_mask	Optional[Tensor]	Mask of obs (1 = observed)	None
guidance	float	Classifier-free guidance scale (1 = off)	1.0

Returns:

Type	Description
Tensor	Generated ECG samples [B, C, L]

Source code in src/ecgen/families/vqvae/prior_lightning.py

@torch.no_grad()
def sample(
    self,
    n_samples: int = 16,
    seq_length: int = 5000,
    temperature: float = 1.0,
    cond: Optional[Mapping[str, Any]] = None,
    obs: Optional[Tensor] = None,
    obs_mask: Optional[Tensor] = None,
    guidance: float = 1.0,
    **kwargs: Any,
) -> Tensor:
    """
    Generate ECG samples by sampling from the prior and decoding.

    Args:
        n_samples: Number of samples to generate
        seq_length: Length of ECG sequence
        temperature: Sampling temperature
        cond: Raw condition values (conditional priors)
        obs: Partially observed ECG (dense conditions)
        obs_mask: Mask of ``obs`` (1 = observed)
        guidance: Classifier-free guidance scale (1 = off)

    Returns:
        Generated ECG samples [B, C, L]
    """
    if self.vqvae is None:
        raise RuntimeError("VQ-VAE not loaded. Cannot decode samples.")

    self.prior.eval()
    device = next(self.prior.parameters()).device

    latent_length = self.vqvae.latent_length(seq_length)

    kwargs_c: dict[str, Any] = {}
    c = self.encode_condition(n_samples, cond, obs, obs_mask)
    if c is not None:
        kwargs_c = {
            "c": c.vector,
            "dense": self._latent_dense(c.dense, latent_length),
        }
        if guidance != 1.0:
            u = self.encode_condition(n_samples, cond, obs, obs_mask, drop="all")
            kwargs_c["c_uncond"] = u.vector
            kwargs_c["dense_uncond"] = self._latent_dense(u.dense, latent_length)
            kwargs_c["guidance"] = guidance

    # Sample discrete codes from prior
    indices = self.prior.sample(
        batch_size=n_samples,
        latent_length=latent_length,
        temperature=temperature,
        device=device,
        **kwargs_c,
    )

    # Decode to ECG
    samples = self.vqvae.decode_from_indices(indices)

    return samples

ecgen.families.gan.pulse2pulse

WaveGANGenerator

Bases: Module

Pulse2Pulse generator from the original WaveGAN implementation, configured for 8‑lead ECG.

Source code in src/ecgen/families/gan/pulse2pulse.py

class WaveGANGenerator(nn.Module):
    """
    Pulse2Pulse generator from the original WaveGAN implementation, configured for 8‑lead ECG.
    """

    def __init__(
        self,
        model_size: int = 50,
        num_channels: int = 8,
        post_proc_filt_len: int = 0,
        verbose: bool = False,
        upsample: bool = True,
        extra_in_channels: int = 0,
        cond_dim: int = 0,
        output_scale: float = 1.0,
    ) -> None:
        super().__init__()
        self.model_size = model_size
        self.num_channels = num_channels
        self.output_scale = output_scale
        self.post_proc_filt_len = post_proc_filt_len
        self.verbose = verbose

        stride = 1 if upsample else 4
        upsample_val: Optional[int] = 5 if upsample else None

        self.deconv_1 = Transpose1dLayer(
            5 * model_size, 5 * model_size, 25, stride, upsample=upsample_val
        )
        self.deconv_2 = Transpose1dLayerMultiInput(
            5 * model_size * 2, 3 * model_size, 25, stride, upsample=upsample_val
        )
        self.deconv_3 = Transpose1dLayerMultiInput(
            3 * model_size * 2, model_size, 25, stride, upsample=upsample_val
        )
        self.deconv_5 = Transpose1dLayerMultiInput(
            model_size * 2, model_size // 2, 25, stride, upsample=2
        )
        self.deconv_6 = Transpose1dLayerMultiInput(
            model_size, model_size // 5, 25, stride, upsample=upsample_val
        )
        self.deconv_7 = Transpose1dLayer(
            model_size // 5, num_channels, 25, stride, upsample=2
        )

        # extra_in_channels: dense conditions concatenated to the input noise
        self.conv_1 = nn.Conv1d(
            num_channels + extra_in_channels,
            model_size // 5,
            25,
            stride=2,
            padding=25 // 2,
        )
        self.conv_2 = nn.Conv1d(
            model_size // 5, model_size // 2, 25, stride=5, padding=25 // 2
        )
        self.conv_3 = nn.Conv1d(
            model_size // 2, model_size, 25, stride=2, padding=25 // 2
        )
        self.conv_4 = nn.Conv1d(
            model_size, model_size * 3, 25, stride=5, padding=25 // 2
        )
        self.conv_5 = nn.Conv1d(
            model_size * 3, model_size * 5, 25, stride=5, padding=25 // 2
        )
        self.conv_6 = nn.Conv1d(
            model_size * 5, model_size * 5, 25, stride=5, padding=25 // 2
        )

        if post_proc_filt_len and post_proc_filt_len > 0:
            self.post_proc_filter = nn.Conv1d(
                num_channels, num_channels, post_proc_filt_len
            )
        else:
            self.post_proc_filter = None

        for m in self.modules():
            if isinstance(m, (nn.ConvTranspose1d, nn.Linear)):
                nn.init.kaiming_normal_(m.weight.data)

        # cond_dim > 0: FiLM on every decoder stage (created after the init loop
        # so they stay zero-initialised)
        self.films = nn.ModuleList()
        if cond_dim > 0:
            for ch in (
                5 * model_size,
                3 * model_size,
                model_size,
                model_size // 2,
                model_size // 5,
            ):
                self.films.append(FiLM(cond_dim, ch))

    def _film(self, i: int, x: Tensor, c: Optional[Tensor]) -> Tensor:
        return self.films[i](x, c) if len(self.films) else x

    def forward(
        self, x: Tensor, c: Optional[Tensor] = None, dense: Optional[Tensor] = None
    ) -> Tensor:
        if dense is not None:
            x = torch.cat([x, dense], dim=1)
        conv_1_out = F.leaky_relu(self.conv_1(x))
        conv_2_out = F.leaky_relu(self.conv_2(conv_1_out))
        conv_3_out = F.leaky_relu(self.conv_3(conv_2_out))
        conv_4_out = F.leaky_relu(self.conv_4(conv_3_out))
        conv_5_out = F.leaky_relu(self.conv_5(conv_4_out))
        x = F.leaky_relu(self.conv_6(conv_5_out))

        x = F.relu(self._film(0, self.deconv_1(x), c))
        x = F.relu(self._film(1, self.deconv_2(x, conv_5_out), c))
        x = F.relu(self._film(2, self.deconv_3(x, conv_4_out), c))
        x = F.relu(self._film(3, self.deconv_5(x, conv_3_out), c))
        x = F.relu(self._film(4, self.deconv_6(x, conv_2_out), c))
        x = torch.tanh(self.deconv_7(x))
        if self.output_scale != 1.0:
            x = x * self.output_scale

        if self.post_proc_filter is not None:
            x = self.post_proc_filter(x)

        return x

PhaseShuffle

Bases: Module

Phase shuffling as in WaveGAN: random circular shifts with reflection padding.

Source code in src/ecgen/families/gan/pulse2pulse.py

class PhaseShuffle(nn.Module):
    """
    Phase shuffling as in WaveGAN: random circular shifts with reflection padding.
    """

    def __init__(self, shift_factor: int) -> None:
        super().__init__()
        self.shift_factor = shift_factor

    def forward(self, x: Tensor) -> Tensor:
        if self.shift_factor == 0:
            return x

        k_list = (
            torch.randint(
                low=0,
                high=2 * self.shift_factor + 1,
                size=(x.shape[0],),
                device=x.device,
            )
            - self.shift_factor
        )

        k_map: dict[int, list[int]] = {}
        for idx, k in enumerate(k_list.tolist()):
            k_map.setdefault(int(k), []).append(idx)

        x_shuffle = x.clone()
        for k, idxs in k_map.items():
            if k > 0:
                x_shuffle[idxs] = F.pad(x[idxs][..., :-k], (k, 0), mode="reflect")
            else:
                x_shuffle[idxs] = F.pad(x[idxs][..., -k:], (0, -k), mode="reflect")

        return x_shuffle

WaveGANDiscriminator

Bases: Module

Pulse2Pulse discriminator with dynamic final linear layer to accommodate sequence length.

Source code in src/ecgen/families/gan/pulse2pulse.py

class WaveGANDiscriminator(nn.Module):
    """
    Pulse2Pulse discriminator with dynamic final linear layer to accommodate sequence length.
    """

    def __init__(
        self,
        model_size: int = 64,
        num_channels: int = 8,
        shift_factor: int = 2,
        alpha: float = 0.2,
        verbose: bool = False,
        extra_in_channels: int = 0,
        cond_dim: int = 0,
    ) -> None:
        super().__init__()
        self.model_size = model_size
        self.num_channels = num_channels
        self.alpha = alpha
        self.verbose = verbose

        # extra_in_channels: dense conditions seen next to the signal (paired D)
        self.conv1 = nn.Conv1d(
            num_channels + extra_in_channels, model_size, 25, stride=2, padding=11
        )
        self.conv2 = nn.Conv1d(model_size, 2 * model_size, 25, stride=2, padding=11)
        self.conv3 = nn.Conv1d(2 * model_size, 5 * model_size, 25, stride=2, padding=11)
        self.conv4 = nn.Conv1d(
            5 * model_size, 10 * model_size, 25, stride=2, padding=11
        )
        self.conv5 = nn.Conv1d(
            10 * model_size, 20 * model_size, 25, stride=4, padding=11
        )
        self.conv6 = nn.Conv1d(
            20 * model_size, 25 * model_size, 25, stride=4, padding=11
        )
        self.conv7 = nn.Conv1d(
            25 * model_size, 100 * model_size, 25, stride=4, padding=11
        )

        self.ps1 = PhaseShuffle(shift_factor)
        self.ps2 = PhaseShuffle(shift_factor)
        self.ps3 = PhaseShuffle(shift_factor)
        self.ps4 = PhaseShuffle(shift_factor)
        self.ps5 = PhaseShuffle(shift_factor)
        self.ps6 = PhaseShuffle(shift_factor)

        self.fc1: Optional[nn.Linear] = None

        for m in self.modules():
            if isinstance(m, (nn.Conv1d, nn.Linear)) and m is not self.fc1:
                nn.init.kaiming_normal_(m.weight.data)

        # Projection discriminator (Miyato & Koyama, 2018): adds
        # <embed(c), mean_t(features)> to the critic output
        self.projection = (
            nn.Linear(cond_dim, 100 * model_size, bias=False) if cond_dim > 0 else None
        )

    def _ensure_fc(self, x: Tensor) -> None:
        flat_dim = x.size(1) * x.size(2)
        if self.fc1 is None or getattr(self.fc1, "in_features", None) != flat_dim:
            self.fc1 = nn.Linear(flat_dim, 1, device=x.device)
            nn.init.kaiming_normal_(self.fc1.weight.data)

    def features(self, x: Tensor) -> Tensor:
        """Convolutional features ``[B, 100 * model_size, L']`` (input incl. dense channels)."""
        x = F.leaky_relu(self.conv1(x), negative_slope=self.alpha)
        x = self.ps1(x)
        x = F.leaky_relu(self.conv2(x), negative_slope=self.alpha)
        x = self.ps2(x)
        x = F.leaky_relu(self.conv3(x), negative_slope=self.alpha)
        x = self.ps3(x)
        x = F.leaky_relu(self.conv4(x), negative_slope=self.alpha)
        x = self.ps4(x)
        x = F.leaky_relu(self.conv5(x), negative_slope=self.alpha)
        x = self.ps5(x)
        x = F.leaky_relu(self.conv6(x), negative_slope=self.alpha)
        x = self.ps6(x)
        x = F.leaky_relu(self.conv7(x), negative_slope=self.alpha)
        return x

    def forward(
        self, x: Tensor, c: Optional[Tensor] = None, dense: Optional[Tensor] = None
    ) -> Tensor:
        if dense is not None:
            x = torch.cat([x, dense], dim=1)
        features = self.features(x)
        x = features

        x = x.view(x.size(0), -1, 1)
        self._ensure_fc(x)
        assert self.fc1 is not None
        x = x.view(x.size(0), -1)
        out = self.fc1(x)
        if self.projection is not None and c is not None:
            out = out + (self.projection(c) * features.mean(dim=-1)).sum(
                dim=1, keepdim=True
            )
        return out

features

features(x: Tensor) -> Tensor

Convolutional features [B, 100 * model_size, L'] (input incl. dense channels).

Source code in src/ecgen/families/gan/pulse2pulse.py

def features(self, x: Tensor) -> Tensor:
    """Convolutional features ``[B, 100 * model_size, L']`` (input incl. dense channels)."""
    x = F.leaky_relu(self.conv1(x), negative_slope=self.alpha)
    x = self.ps1(x)
    x = F.leaky_relu(self.conv2(x), negative_slope=self.alpha)
    x = self.ps2(x)
    x = F.leaky_relu(self.conv3(x), negative_slope=self.alpha)
    x = self.ps3(x)
    x = F.leaky_relu(self.conv4(x), negative_slope=self.alpha)
    x = self.ps4(x)
    x = F.leaky_relu(self.conv5(x), negative_slope=self.alpha)
    x = self.ps5(x)
    x = F.leaky_relu(self.conv6(x), negative_slope=self.alpha)
    x = self.ps6(x)
    x = F.leaky_relu(self.conv7(x), negative_slope=self.alpha)
    return x

Pulse2PulseGAN

Bases: GenerativeModule

PyTorch Lightning wrapper around the Pulse2Pulse WaveGAN generator + discriminator.

With config.conditioning it is a conditional WGAN-GP:

• global conditions (vector) modulate the generator's decoder with FiLM and enter the critic through a projection head; the critic has its own copy of the conditioner (trained by the critic's optimizer, without dropout);
• time-aligned conditions (dense, e.g. masked_signal) are concatenated to the generator's input noise and to the critic's input, which makes it a paired (pix2pix-style) translator, e.g. 1 lead -> 8 leads; set l1_weight (e.g. 100) to add an L1 term on the unobserved part.

Source code in src/ecgen/families/gan/pulse2pulse.py

@FAMILIES.register("pulse2pulse")
class Pulse2PulseGAN(GenerativeModule):
    """
    PyTorch Lightning wrapper around the Pulse2Pulse WaveGAN generator + discriminator.

    With ``config.conditioning`` it is a conditional WGAN-GP:

    - global conditions (``vector``) modulate the generator's decoder with FiLM and
      enter the critic through a projection head; the critic has its own copy of
      the conditioner (trained by the critic's optimizer, without dropout);
    - time-aligned conditions (``dense``, e.g. ``masked_signal``) are concatenated
      to the generator's input noise and to the critic's input, which makes it a
      paired (pix2pix-style) translator, e.g. 1 lead -> 8 leads; set ``l1_weight``
      (e.g. 100) to add an L1 term on the unobserved part.
    """

    family = "gan"
    supports = Capabilities(conditional=True)
    accepts = frozenset({"vector", "dense"})

    def __init__(
        self, config: Pulse2PulseConfig | dict[str, Any] | None = None, **kwargs: Any
    ) -> None:
        super().__init__()
        config = resolve_config(Pulse2PulseConfig, config, kwargs)
        self.save_hyperparameters(config.__dict__)
        self.config = config
        if config.l1_on not in ("all", "unobserved"):
            raise ValueError(
                f"l1_on must be 'all' or 'unobserved', got {config.l1_on!r}"
            )

        self.init_conditioning(config.conditioning)
        cond_dim = dense_channels = 0
        self.d_conditioner = None
        if self.conditioner is not None:
            cond_dim = (
                self.conditioner.vector_dim if "vector" in self.conditioner.slots else 0
            )
            dense_channels = self.conditioner.dense_channels
            d_spec = {k: v for k, v in config.conditioning.items() if k != "dropout"}
            self.d_conditioner = build_conditioner(d_spec)

        self.netG = WaveGANGenerator(
            model_size=config.model_size,
            num_channels=config.num_channels,
            extra_in_channels=dense_channels,
            cond_dim=cond_dim,
            output_scale=config.output_scale,
        )
        self.netD = WaveGANDiscriminator(
            model_size=config.model_size,
            num_channels=config.num_channels,
            extra_in_channels=dense_channels,
            cond_dim=cond_dim,
        )
        # Build netD.fc1 now: created lazily it would be missing from the
        # critic's optimizer (configure_optimizers runs before the first forward).
        with torch.no_grad():
            dummy = torch.zeros(
                1, config.num_channels + dense_channels, config.seq_length
            )
            self.netD._ensure_fc(self.netD.features(dummy).reshape(1, -1, 1))

        self.automatic_optimization = False

        self._val_real_sample: Tensor | None = None
        self._val_fake_sample: Tensor | None = None

    def forward(self, noise: Tensor) -> Tensor:
        return self.netG(noise)

    @property
    def fixed_seq_length(self) -> Optional[int]:
        return self.config.seq_length

    def on_load_checkpoint(self, checkpoint: dict[str, Any]) -> None:
        super().on_load_checkpoint(checkpoint)
        # Older checkpoints may hold an fc1 of another size (lazy creation);
        # rebuild it at the saved size so the weights load.
        fc_weight = checkpoint.get("state_dict", {}).get("netD.fc1.weight")
        if fc_weight is not None and (
            self.netD.fc1 is None or self.netD.fc1.in_features != fc_weight.shape[1]
        ):
            self.netD.fc1 = nn.Linear(fc_weight.shape[1], 1)

    def _sample_noise(self, batch_size: int, device: torch.device) -> Tensor:
        return torch.empty(
            batch_size,
            self.config.num_channels,
            self.config.seq_length,
            device=device,
        ).uniform_(-1, 1)

    # ------------------------------------------------------------ conditions

    def _conditions(self, batch: Any):
        """(generator condition, critic condition) for a batch; (None, None) if unconditional."""
        if self.conditioner is None:
            return None, None
        g = self.condition(batch)
        obs = batch.signal if batch.obs_mask is not None else None
        d = self.d_conditioner(
            batch.cond,
            batch.cond_missing,
            obs=obs,
            obs_mask=batch.obs_mask,
            batch_size=len(batch),
        )
        return g, d

    @staticmethod
    def _slots(c) -> tuple[Optional[Tensor], Optional[Tensor]]:
        return (None, None) if c is None else (c.vector, c.dense)

    def _critic(self, x: Tensor, d_cond) -> Tensor:
        vec, dense = self._slots(d_cond)
        return self.netD(x, vec, dense)

    def _generate(self, b: int, device: torch.device, g_cond) -> Tensor:
        vec, dense = self._slots(g_cond)
        return self.netG(self._sample_noise(b, device), vec, dense)

    def _l1(self, fake: Tensor, real: Tensor, batch: Any) -> Tensor:
        obs_mask = getattr(batch, "obs_mask", None)
        if self.config.l1_on == "unobserved" and obs_mask is not None:
            w = 1.0 - obs_mask.to(real.dtype)
            return ((fake - real).abs() * w).sum() / w.sum().clamp_min(1.0)
        return (fake - real).abs().mean()

    # -------------------------------------------------------------- training

    def training_step(self, batch: Any, batch_idx: int) -> Tensor:
        opt_g, opt_d = self.optimizers()

        real_ecgs = self.get_signal(batch)
        device = real_ecgs.device
        b_size = real_ecgs.size(0)
        g_cond, d_cond = self._conditions(batch)

        # --- Update D ---
        opt_d.zero_grad()
        fake = self._generate(b_size, device, g_cond).detach()

        d_real = self._critic(real_ecgs, d_cond).mean()
        d_fake = self._critic(fake, d_cond).mean()
        gp = calc_gradient_penalty(
            lambda x: self._critic(x, d_cond), real_ecgs, fake, self.config.lmbda
        )

        d_loss = d_fake - d_real + gp
        self.manual_backward(d_loss)
        opt_d.step()

        d_wasserstein = d_real - d_fake

        # --- Update G every n_critic steps ---
        g_loss = torch.tensor(0.0, device=device)
        if (batch_idx + 1) % self.config.n_critic == 0:
            opt_g.zero_grad()
            if g_cond is not None:  # fresh condition dropout for the G step
                g_cond, d_cond = self._conditions(batch)
            fake = self._generate(b_size, device, g_cond)
            g_adv = self._critic(fake, d_cond).mean()
            g_loss = -g_adv
            if self.config.l1_weight > 0:
                l1 = self._l1(fake, real_ecgs, batch)
                g_loss = g_loss + self.config.l1_weight * l1
                self.log("train/l1", l1, on_step=True, on_epoch=True)
            self.manual_backward(g_loss)
            opt_g.step()

        self.log("train/d_loss", d_loss, on_step=True, on_epoch=True, prog_bar=True)
        self.log(
            "train/d_wasserstein",
            d_wasserstein,
            on_step=True,
            on_epoch=True,
            prog_bar=False,
        )
        self.log("train/g_loss", g_loss, on_step=True, on_epoch=True, prog_bar=True)

        return d_loss

    def validation_step(self, batch: Any, batch_idx: int) -> Tensor:
        real_ecgs = self.get_signal(batch)
        device = real_ecgs.device
        b_size = real_ecgs.size(0)
        g_cond, d_cond = self._conditions(batch)

        fake = self._generate(b_size, device, g_cond)

        d_real = self._critic(real_ecgs, d_cond).mean()
        d_fake = self._critic(fake, d_cond).mean()
        d_wasserstein = d_real - d_fake

        self.log(
            "val/d_wasserstein",
            d_wasserstein,
            on_step=False,
            on_epoch=True,
            prog_bar=True,
        )
        self.log(
            "val_loss", -d_wasserstein, on_step=False, on_epoch=True, prog_bar=True
        )
        if g_cond is not None and self.config.l1_weight > 0:
            self.log(
                "val/l1", self._l1(fake, real_ecgs, batch), on_step=False, on_epoch=True
            )

        if batch_idx == 0:
            self._val_real_sample = real_ecgs[0].detach().cpu()
            self._val_fake_sample = fake[0].detach().cpu()

        return d_wasserstein

    # -------------------------------------------------------------- sampling

    @torch.no_grad()
    def generate_samples(self, n_samples: int = 16) -> Tensor:
        """Generate ECG samples from random noise (unconditional models)."""
        self.netG.eval()
        device = next(self.netG.parameters()).device
        noise = self._sample_noise(n_samples, device)
        fake = self.netG(noise)
        return fake

    @torch.no_grad()
    def sample(
        self,
        n_samples: int = 16,
        cond: Optional[dict[str, Any]] = None,
        obs: Optional[Tensor] = None,
        obs_mask: Optional[Tensor] = None,
        paste_observed: bool = True,
        **kwargs: Any,
    ) -> Tensor:
        """
        Generate ECGs; with conditioning, given raw ``cond`` values and/or a
        partially observed ``obs`` (``[n, C, T]``) with ``obs_mask``.
        """
        if self.conditioner is None:
            if cond or obs is not None:
                raise ValueError(
                    "This Pulse2Pulse model was trained without conditioning"
                )
            return self.generate_samples(n_samples)
        self.netG.eval()
        c = self.encode_condition(n_samples, cond, obs, obs_mask)
        fake = self._generate(n_samples, self.device, c)
        if paste_observed and obs is not None:
            m = torch.ones_like(obs) if obs_mask is None else obs_mask
            m, obs = m.to(fake), obs.to(fake)
            fake = obs * m + fake * (1 - m)
        return fake

    def configure_optimizers(self) -> Any:
        g_params = list(self.netG.parameters())
        d_params = list(self.netD.parameters())
        if self.conditioner is not None:
            g_params += list(self.conditioner.parameters())
            d_params += list(self.d_conditioner.parameters())
        opt_g = torch.optim.Adam(
            g_params, lr=self.config.lr, betas=(self.config.b1, self.config.b2)
        )
        opt_d = torch.optim.Adam(
            d_params, lr=self.config.lr, betas=(self.config.b1, self.config.b2)
        )
        return [opt_g, opt_d], []

generate_samples

generate_samples(n_samples: int = 16) -> Tensor

Generate ECG samples from random noise (unconditional models).

Source code in src/ecgen/families/gan/pulse2pulse.py

@torch.no_grad()
def generate_samples(self, n_samples: int = 16) -> Tensor:
    """Generate ECG samples from random noise (unconditional models)."""
    self.netG.eval()
    device = next(self.netG.parameters()).device
    noise = self._sample_noise(n_samples, device)
    fake = self.netG(noise)
    return fake

sample

sample(n_samples: int = 16, cond: Optional[dict[str, Any]] = None, obs: Optional[Tensor] = None, obs_mask: Optional[Tensor] = None, paste_observed: bool = True, **kwargs: Any) -> Tensor

Generate ECGs; with conditioning, given raw cond values and/or a partially observed obs ([n, C, T]) with obs_mask.

Source code in src/ecgen/families/gan/pulse2pulse.py

@torch.no_grad()
def sample(
    self,
    n_samples: int = 16,
    cond: Optional[dict[str, Any]] = None,
    obs: Optional[Tensor] = None,
    obs_mask: Optional[Tensor] = None,
    paste_observed: bool = True,
    **kwargs: Any,
) -> Tensor:
    """
    Generate ECGs; with conditioning, given raw ``cond`` values and/or a
    partially observed ``obs`` (``[n, C, T]``) with ``obs_mask``.
    """
    if self.conditioner is None:
        if cond or obs is not None:
            raise ValueError(
                "This Pulse2Pulse model was trained without conditioning"
            )
        return self.generate_samples(n_samples)
    self.netG.eval()
    c = self.encode_condition(n_samples, cond, obs, obs_mask)
    fake = self._generate(n_samples, self.device, c)
    if paste_observed and obs is not None:
        m = torch.ones_like(obs) if obs_mask is None else obs_mask
        m, obs = m.to(fake), obs.to(fake)
        fake = obs * m + fake * (1 - m)
    return fake

calc_gradient_penalty

calc_gradient_penalty(netD: Module, real_data: Tensor, fake_data: Tensor, lmbda: float) -> Tensor

WGAN‑GP gradient penalty.

Source code in src/ecgen/families/gan/pulse2pulse.py

def calc_gradient_penalty(
    netD: nn.Module,
    real_data: Tensor,
    fake_data: Tensor,
    lmbda: float,
) -> Tensor:
    """
    WGAN‑GP gradient penalty.
    """
    batch_size = real_data.size(0)
    alpha = torch.rand(batch_size, 1, 1, device=real_data.device)
    alpha = alpha.expand_as(real_data)
    interpolates = alpha * real_data + (1 - alpha) * fake_data
    interpolates.requires_grad_(True)
    disc_interpolates = netD(interpolates)
    gradients = torch.autograd.grad(
        outputs=disc_interpolates,
        inputs=interpolates,
        grad_outputs=torch.ones_like(disc_interpolates, device=real_data.device),
        create_graph=True,
        retain_graph=True,
        only_inputs=True,
    )[0]
    gradients = gradients.view(batch_size, -1)
    gradient_penalty = ((gradients.norm(2, dim=1) - 1) ** 2).mean() * lmbda
    return gradient_penalty

ecgen.families.continuous

One LightningModule for every continuous-time generative method (diffusion, flow matching): the method is a :class:~ecgen.processes.Process, the network a registered backbone, the conditions a :class:~ecgen.conditioning.CompositeConditioner.

ContinuousConfig dataclass

Attributes:

Name	Type	Description
in_channels	int	Leads.
process	Dict[str, Any]	{name, params} from ecgen.registry.PROCESSES (ddpm, ...).
backbone	Dict[str, Any]	{name, params} from ecgen.registry.BACKBONES (unet1d, ...); channel and condition sizes are filled in.
conditioning	Optional[Dict[str, Any]]	See :mod:ecgen.conditioning (None = unconditional).
lr,	(weight_decay, b1, b2)	AdamW.
warmup_steps	int	Linear learning-rate warm-up.
ema_decay	float	EMA of the weights used for sampling (0 = off).
loss_on	str	"all" or "unobserved" (with data.masks): loss only where obs_mask is 0.
clean_observed	bool	SSSD-style partial-ECG training: the observed entries of the noisy input are replaced by the clean signal (and again at every sampling step).
sample_steps	int	Default sampling steps.
guidance	float	Default classifier-free guidance scale (1 = off).

Source code in src/ecgen/families/continuous/module.py

@dataclass
class ContinuousConfig:
    """
    Attributes:
        in_channels: Leads.
        process: ``{name, params}`` from ``ecgen.registry.PROCESSES`` (``ddpm``, ...).
        backbone: ``{name, params}`` from ``ecgen.registry.BACKBONES`` (``unet1d``, ...);
            channel and condition sizes are filled in.
        conditioning: See :mod:`ecgen.conditioning` (``None`` = unconditional).
        lr, weight_decay, b1, b2: AdamW.
        warmup_steps: Linear learning-rate warm-up.
        ema_decay: EMA of the weights used for sampling (0 = off).
        loss_on: ``"all"`` or ``"unobserved"`` (with ``data.masks``): loss only where
            ``obs_mask`` is 0.
        clean_observed: SSSD-style partial-ECG training: the observed entries of the
            noisy input are replaced by the clean signal (and again at every
            sampling step).
        sample_steps: Default sampling steps.
        guidance: Default classifier-free guidance scale (1 = off).
    """

    in_channels: int = 12
    process: Dict[str, Any] = field(default_factory=lambda: {"name": "ddpm"})
    backbone: Dict[str, Any] = field(default_factory=lambda: {"name": "unet1d"})
    conditioning: Optional[Dict[str, Any]] = None
    lr: float = 2e-4
    weight_decay: float = 0.0
    b1: float = 0.9
    b2: float = 0.999
    warmup_steps: int = 500
    ema_decay: float = 0.999
    loss_on: str = "all"
    clean_observed: bool = False
    sample_steps: int = 50
    guidance: float = 1.0

ContinuousGenerativeModule

Bases: GenerativeModule

Diffusion / flow-matching family. Registered as diffusion (default process ddpm) and flow_matching (default process rectified_flow).

Training: draw noise, let the process build (x_t, t, target, weight), regress the backbone output on the target. Sampling: the process integrates the backbone from noise, with optional classifier-free guidance, and partial-ECG conditioning either learned (masked_signal encoder, clean_observed) or training-free (inpaint="replace" | "repaint").

Source code in src/ecgen/families/continuous/module.py

class ContinuousGenerativeModule(GenerativeModule):
    """
    Diffusion / flow-matching family. Registered as ``diffusion`` (default
    process ``ddpm``) and ``flow_matching`` (default process ``rectified_flow``).

    Training: draw noise, let the process build ``(x_t, t, target, weight)``,
    regress the backbone output on the target. Sampling: the process integrates
    the backbone from noise, with optional classifier-free guidance, and
    partial-ECG conditioning either learned (``masked_signal`` encoder,
    ``clean_observed``) or training-free (``inpaint="replace" | "repaint"``).
    """

    family = "continuous"
    supports = Capabilities(conditional=True, cfg=True, inpainting=True)
    accepts = frozenset({"vector", "tokens", "dense"})
    config_cls: type = ContinuousConfig

    def __init__(
        self, config: ContinuousConfig | dict[str, Any] | None = None, **kwargs: Any
    ) -> None:
        super().__init__()
        config = resolve_config(self.config_cls, config, kwargs)
        if config.loss_on not in ("all", "unobserved"):
            raise ValueError(
                f"loss_on must be 'all' or 'unobserved', got {config.loss_on!r}"
            )
        self.save_hyperparameters(config.__dict__)
        self.config = config

        self.init_conditioning(config.conditioning)
        sizes = {"cond_dim": 0, "context_dim": 0, "dense_channels": 0}
        if self.conditioner is not None:
            slots = self.conditioner.slots
            sizes["cond_dim"] = self.conditioner.vector_dim if "vector" in slots else 0
            sizes["context_dim"] = (
                self.conditioner.token_dim if "tokens" in slots else 0
            )
            sizes["dense_channels"] = self.conditioner.dense_channels
            backbone_cls = BACKBONES.resolve(config.backbone)
            extra = slots - set(getattr(backbone_cls, "accepts", self.accepts))
            if extra:
                raise ValueError(
                    f"Backbone {backbone_cls.__name__} cannot consume condition slots {sorted(extra)}"
                )
        self.process = PROCESSES.build(config.process)
        self.backbone = BACKBONES.build(
            config.backbone, in_channels=self._model_channels(config), **sizes
        )
        self.ema = (
            EMA([self.backbone, self.conditioner], config.ema_decay)
            if config.ema_decay > 0
            else None
        )

    # ------------------------------------------------- model space (hooks)

    def _model_channels(self, config: ContinuousConfig) -> int:
        """Channels of the space the process runs in (latent subclasses override)."""
        return config.in_channels

    def to_model_space(self, x: Tensor) -> Tensor:
        """Signal ``[B, C, T]`` → the space the process runs in (identity here)."""
        return x

    def from_model_space(self, z: Tensor) -> Tensor:
        """Inverse of :meth:`to_model_space`."""
        return z

    def model_space_shape(self, n: int, seq_length: int) -> tuple[int, ...]:
        return (n, self.config.in_channels, seq_length)

    def adapt_condition(
        self, c: Optional[ConditionOutput], length: int
    ) -> Optional[ConditionOutput]:
        """Resample dense conditions to the model-space length."""
        if c is None or c.dense is None or c.dense.shape[-1] == length:
            return c
        dense = torch.nn.functional.adaptive_avg_pool1d(c.dense, length)
        return ConditionOutput(c.vector, c.tokens, c.tokens_mask, dense)

    # ------------------------------------------------------------- training

    def _loss(self, batch: Any) -> Tensor:
        signal = self.get_signal(batch)
        x0 = self.to_model_space(signal)
        cond = self.adapt_condition(self.condition(batch), x0.shape[-1])
        obs_mask = getattr(batch, "obs_mask", None)
        if x0.shape != signal.shape:
            obs_mask = None  # masks live in signal space
        pair = self.process.training_pair(x0, torch.randn_like(x0))
        xt = pair.xt
        if self.config.clean_observed and obs_mask is not None:
            m = obs_mask.to(x0.dtype)
            xt = m * x0 + (1 - m) * xt
        out = self.backbone(xt, pair.t, cond)
        err = (out - pair.target) ** 2
        if self.config.loss_on == "unobserved" and obs_mask is not None:
            w = 1.0 - obs_mask.to(err.dtype)
            per_sample = (err * w).flatten(1).sum(1) / w.flatten(1).sum(1).clamp_min(
                1.0
            )
        else:
            per_sample = err.flatten(1).mean(1)
        return (pair.weight * per_sample).mean()

    def training_step(self, batch: Any, batch_idx: int) -> Tensor:
        loss = self._loss(batch)
        self.log("train/loss", loss, on_step=True, on_epoch=True, prog_bar=True)
        return loss

    def validation_step(self, batch: Any, batch_idx: int) -> Tensor:
        loss = self._loss(batch)
        self.log("val/loss", loss, on_step=False, on_epoch=True, prog_bar=True)
        self.log("val_loss", loss, on_step=False, on_epoch=True, prog_bar=True)
        return loss

    def on_train_batch_end(self, outputs: Any, batch: Any, batch_idx: int) -> None:
        if self.ema is not None:
            self.ema.update()

    def configure_optimizers(self) -> Any:
        params = [p for p in self.parameters() if p.requires_grad]
        opt = torch.optim.AdamW(
            params,
            lr=self.config.lr,
            betas=(self.config.b1, self.config.b2),
            weight_decay=self.config.weight_decay,
        )
        warmup = max(1, self.config.warmup_steps)
        sched = torch.optim.lr_scheduler.LambdaLR(
            opt, lambda step: min(1.0, (step + 1) / warmup)
        )
        return {
            "optimizer": opt,
            "lr_scheduler": {"scheduler": sched, "interval": "step"},
        }

    # ----------------------------------------------------------- checkpoint

    def on_save_checkpoint(self, checkpoint: Dict[str, Any]) -> None:
        super().on_save_checkpoint(checkpoint)
        if self.ema is not None:
            checkpoint["ema"] = self.ema.state_dict()

    def on_load_checkpoint(self, checkpoint: Dict[str, Any]) -> None:
        super().on_load_checkpoint(checkpoint)
        if self.ema is not None and "ema" in checkpoint:
            self.ema.load_state_dict(checkpoint["ema"])

    # ------------------------------------------------------------- sampling

    def _model_fn(self, c: Optional[ConditionOutput], length: int):
        c = self.adapt_condition(c, length)
        return lambda x, t: self.backbone(x, t, c)

    @torch.no_grad()
    def sample(
        self,
        n_samples: int = 16,
        seq_length: Optional[int] = None,
        cond: Optional[Mapping[str, Any]] = None,
        obs: Optional[Tensor] = None,
        obs_mask: Optional[Tensor] = None,
        guidance: Optional[float] = None,
        steps: Optional[int] = None,
        inpaint: Optional[str] = None,
        jumps: int = 5,
        paste_observed: bool = True,
        use_ema: bool = True,
        seed: Optional[int] = None,
        **process_kwargs: Any,
    ) -> Tensor:
        """
        Generate ``[n, C, T]`` ECGs.

        Args:
            n_samples: Number of ECGs.
            seq_length: Length; defaults to the training data's.
            cond: Raw condition values (see :meth:`prepare_condition`).
            obs: Partially observed ECG. Used by a ``masked_signal`` condition,
                ``clean_observed`` models, and ``inpaint``.
            obs_mask: Mask of ``obs`` (1 = observed).
            guidance: Classifier-free guidance scale (default ``config.guidance``).
            steps: Sampling steps (default ``config.sample_steps``).
            inpaint: Training-free inpainting for models without a learned observation
                condition: ``"replace"`` or ``"repaint"`` (needs ``obs``).
            jumps: RePaint resampling count.
            paste_observed: Copy the observed part of ``obs`` into the output.
            use_ema: Sample with the EMA weights.
            seed: Seed for the sampling noise.
            process_kwargs: Passed to the process sampler (e.g. ``eta`` for DDIM,
                ``solver`` for flows).
        """
        self.eval()
        seq_length = seq_length or self.data_info.get("seq_length")
        if seq_length is None:
            raise ValueError("seq_length is unknown; pass seq_length=")
        guidance = self.config.guidance if guidance is None else guidance
        steps = steps or self.config.sample_steps
        generator = (
            torch.Generator(device=self.device).manual_seed(seed)
            if seed is not None
            else None
        )
        if obs is not None:
            obs = obs.to(self.device)
            obs_mask = (
                torch.ones_like(obs)
                if obs_mask is None
                else obs_mask.to(self.device, obs.dtype)
            )

        shape = self.model_space_shape(n_samples, seq_length)
        latent = shape[1:] != (self.config.in_channels, seq_length)
        if latent and (inpaint is not None or self.config.clean_observed):
            raise ValueError("Training-free inpainting works in signal space only")
        inpainting = None
        if inpaint is not None:
            if obs is None:
                raise ValueError("inpaint needs obs (and obs_mask)")
            inpainting = Inpainting(obs, obs_mask, inpaint, jumps)
        elif self.config.clean_observed and obs is not None:
            inpainting = Inpainting(obs, obs_mask, "clean")

        scope = (
            self.ema.swap() if (use_ema and self.ema is not None) else _nullcontext()
        )
        # obs feeds the conditioner only when it has an observation encoder;
        # otherwise it is used for inpainting / pasting only
        c_obs, c_mask = (
            (obs, obs_mask)
            if (self.conditioner is not None and self.conditioner.uses_obs)
            else (None, None)
        )
        with scope:
            c = self.encode_condition(n_samples, cond, c_obs, c_mask)
            uncond_fn = None
            if c is not None and guidance != 1.0:
                u = self.encode_condition(n_samples, cond, c_obs, c_mask, drop="all")
                uncond_fn = self._model_fn(u, shape[-1])
            model_fn = cfg_model_fn(self._model_fn(c, shape[-1]), uncond_fn, guidance)
            x = self.process.sample(
                model_fn,
                shape,
                steps=steps,
                device=self.device,
                generator=generator,
                inpaint=inpainting,
                **process_kwargs,
            )
            x = self.from_model_space(x)
        if paste_observed and obs is not None:
            x = obs_mask * obs + (1 - obs_mask) * x
        return x

to_model_space

to_model_space(x: Tensor) -> Tensor

Signal [B, C, T] → the space the process runs in (identity here).

Source code in src/ecgen/families/continuous/module.py

def to_model_space(self, x: Tensor) -> Tensor:
    """Signal ``[B, C, T]`` → the space the process runs in (identity here)."""
    return x

from_model_space

from_model_space(z: Tensor) -> Tensor

Inverse of :meth:to_model_space.

Source code in src/ecgen/families/continuous/module.py

def from_model_space(self, z: Tensor) -> Tensor:
    """Inverse of :meth:`to_model_space`."""
    return z

adapt_condition

adapt_condition(c: Optional[ConditionOutput], length: int) -> Optional[ConditionOutput]

Resample dense conditions to the model-space length.

Source code in src/ecgen/families/continuous/module.py

def adapt_condition(
    self, c: Optional[ConditionOutput], length: int
) -> Optional[ConditionOutput]:
    """Resample dense conditions to the model-space length."""
    if c is None or c.dense is None or c.dense.shape[-1] == length:
        return c
    dense = torch.nn.functional.adaptive_avg_pool1d(c.dense, length)
    return ConditionOutput(c.vector, c.tokens, c.tokens_mask, dense)

sample

sample(n_samples: int = 16, seq_length: Optional[int] = None, cond: Optional[Mapping[str, Any]] = None, obs: Optional[Tensor] = None, obs_mask: Optional[Tensor] = None, guidance: Optional[float] = None, steps: Optional[int] = None, inpaint: Optional[str] = None, jumps: int = 5, paste_observed: bool = True, use_ema: bool = True, seed: Optional[int] = None, **process_kwargs: Any) -> Tensor

Generate [n, C, T] ECGs.

Parameters:

Name	Type	Description	Default
n_samples	int	Number of ECGs.	16
seq_length	Optional[int]	Length; defaults to the training data's.	None
cond	Optional[Mapping[str, Any]]	Raw condition values (see :meth:prepare_condition).	None
obs	Optional[Tensor]	Partially observed ECG. Used by a masked_signal condition, clean_observed models, and inpaint.	None
obs_mask	Optional[Tensor]	Mask of obs (1 = observed).	None
guidance	Optional[float]	Classifier-free guidance scale (default config.guidance).	None
steps	Optional[int]	Sampling steps (default config.sample_steps).	None
inpaint	Optional[str]	Training-free inpainting for models without a learned observation condition: "replace" or "repaint" (needs obs).	None
jumps	int	RePaint resampling count.	5
paste_observed	bool	Copy the observed part of obs into the output.	True
use_ema	bool	Sample with the EMA weights.	True
seed	Optional[int]	Seed for the sampling noise.	None
process_kwargs	Any	Passed to the process sampler (e.g. eta for DDIM, solver for flows).	{}

Source code in src/ecgen/families/continuous/module.py

@torch.no_grad()
def sample(
    self,
    n_samples: int = 16,
    seq_length: Optional[int] = None,
    cond: Optional[Mapping[str, Any]] = None,
    obs: Optional[Tensor] = None,
    obs_mask: Optional[Tensor] = None,
    guidance: Optional[float] = None,
    steps: Optional[int] = None,
    inpaint: Optional[str] = None,
    jumps: int = 5,
    paste_observed: bool = True,
    use_ema: bool = True,
    seed: Optional[int] = None,
    **process_kwargs: Any,
) -> Tensor:
    """
    Generate ``[n, C, T]`` ECGs.

    Args:
        n_samples: Number of ECGs.
        seq_length: Length; defaults to the training data's.
        cond: Raw condition values (see :meth:`prepare_condition`).
        obs: Partially observed ECG. Used by a ``masked_signal`` condition,
            ``clean_observed`` models, and ``inpaint``.
        obs_mask: Mask of ``obs`` (1 = observed).
        guidance: Classifier-free guidance scale (default ``config.guidance``).
        steps: Sampling steps (default ``config.sample_steps``).
        inpaint: Training-free inpainting for models without a learned observation
            condition: ``"replace"`` or ``"repaint"`` (needs ``obs``).
        jumps: RePaint resampling count.
        paste_observed: Copy the observed part of ``obs`` into the output.
        use_ema: Sample with the EMA weights.
        seed: Seed for the sampling noise.
        process_kwargs: Passed to the process sampler (e.g. ``eta`` for DDIM,
            ``solver`` for flows).
    """
    self.eval()
    seq_length = seq_length or self.data_info.get("seq_length")
    if seq_length is None:
        raise ValueError("seq_length is unknown; pass seq_length=")
    guidance = self.config.guidance if guidance is None else guidance
    steps = steps or self.config.sample_steps
    generator = (
        torch.Generator(device=self.device).manual_seed(seed)
        if seed is not None
        else None
    )
    if obs is not None:
        obs = obs.to(self.device)
        obs_mask = (
            torch.ones_like(obs)
            if obs_mask is None
            else obs_mask.to(self.device, obs.dtype)
        )

    shape = self.model_space_shape(n_samples, seq_length)
    latent = shape[1:] != (self.config.in_channels, seq_length)
    if latent and (inpaint is not None or self.config.clean_observed):
        raise ValueError("Training-free inpainting works in signal space only")
    inpainting = None
    if inpaint is not None:
        if obs is None:
            raise ValueError("inpaint needs obs (and obs_mask)")
        inpainting = Inpainting(obs, obs_mask, inpaint, jumps)
    elif self.config.clean_observed and obs is not None:
        inpainting = Inpainting(obs, obs_mask, "clean")

    scope = (
        self.ema.swap() if (use_ema and self.ema is not None) else _nullcontext()
    )
    # obs feeds the conditioner only when it has an observation encoder;
    # otherwise it is used for inpainting / pasting only
    c_obs, c_mask = (
        (obs, obs_mask)
        if (self.conditioner is not None and self.conditioner.uses_obs)
        else (None, None)
    )
    with scope:
        c = self.encode_condition(n_samples, cond, c_obs, c_mask)
        uncond_fn = None
        if c is not None and guidance != 1.0:
            u = self.encode_condition(n_samples, cond, c_obs, c_mask, drop="all")
            uncond_fn = self._model_fn(u, shape[-1])
        model_fn = cfg_model_fn(self._model_fn(c, shape[-1]), uncond_fn, guidance)
        x = self.process.sample(
            model_fn,
            shape,
            steps=steps,
            device=self.device,
            generator=generator,
            inpaint=inpainting,
            **process_kwargs,
        )
        x = self.from_model_space(x)
    if paste_observed and obs is not None:
        x = obs_mask * obs + (1 - obs_mask) * x
    return x

DiffusionModule

Bases: ContinuousGenerativeModule

Gaussian diffusion (default process ddpm); see :class:ContinuousGenerativeModule.

Source code in src/ecgen/families/continuous/module.py

@FAMILIES.register("diffusion")
class DiffusionModule(ContinuousGenerativeModule):
    """Gaussian diffusion (default process ``ddpm``); see :class:`ContinuousGenerativeModule`."""

    family = "diffusion"

FlowMatchingConfig dataclass

Bases: ContinuousConfig

:class:ContinuousConfig with the rectified-flow process as default.

Source code in src/ecgen/families/continuous/module.py

@dataclass
class FlowMatchingConfig(ContinuousConfig):
    """:class:`ContinuousConfig` with the rectified-flow process as default."""

    process: Dict[str, Any] = field(default_factory=lambda: {"name": "rectified_flow"})

FlowMatchingModule

Bases: ContinuousGenerativeModule

Flow matching (default process rectified_flow); see :class:ContinuousGenerativeModule.

Source code in src/ecgen/families/continuous/module.py

@FAMILIES.register("flow_matching")
class FlowMatchingModule(ContinuousGenerativeModule):
    """Flow matching (default process ``rectified_flow``); see :class:`ContinuousGenerativeModule`."""

    family = "flow_matching"
    config_cls = FlowMatchingConfig

Latent diffusion / flow matching: the process runs in the latent space of a frozen, pre-trained autoencoder (stage 1), as in latent diffusion models (Rombach et al., 2022).

LatentConfig dataclass

Bases: ContinuousConfig

:class:ContinuousConfig plus the frozen stage-1 autoencoder.

Attributes:

Name	Type	Description
autoencoder_checkpoint	str	Stage-1 checkpoint (vae with latent_type: spatial, or vqvae).
autoencoder_family	str	"vae" or "vqvae".
autoencoder_config	Optional[Dict[str, Any]]	Filled in from the checkpoint, so latent checkpoints restore the autoencoder from their own weights (stage 1 is not needed again).
latent_scale	Optional[float]	Latents are multiplied by this before the process (None = 1 / std of the first training batch's latents, then stored).

Source code in src/ecgen/families/continuous/latent.py

@dataclass
class LatentConfig(ContinuousConfig):
    """
    :class:`ContinuousConfig` plus the frozen stage-1 autoencoder.

    Attributes:
        autoencoder_checkpoint: Stage-1 checkpoint (``vae`` with ``latent_type: spatial``,
            or ``vqvae``).
        autoencoder_family: ``"vae"`` or ``"vqvae"``.
        autoencoder_config: Filled in from the checkpoint, so latent checkpoints
            restore the autoencoder from their own weights (stage 1 is not needed again).
        latent_scale: Latents are multiplied by this before the process (``None`` =
            ``1 / std`` of the first training batch's latents, then stored).
    """

    autoencoder_checkpoint: str = ""
    autoencoder_family: str = "vae"
    autoencoder_config: Optional[Dict[str, Any]] = None
    latent_scale: Optional[float] = None
    backbone: Dict[str, Any] = field(
        default_factory=lambda: {
            "name": "unet1d",
            "params": {"channel_mults": [1, 2, 2]},
        }
    )

LatentContinuousModule

Bases: ContinuousGenerativeModule

Diffusion / flow in a frozen autoencoder's latent space.

• vae stage 1: the encoder's posterior mean is the latent; decode with the decoder.
• vqvae stage 1: the quantized embeddings are the latent; sampled latents are snapped to the codebook before decoding.

Dense conditions are average-pooled to the latent length; observation masks and training-free inpainting are signal-space features and are not used here (a masked_signal condition still works).

Source code in src/ecgen/families/continuous/latent.py

class LatentContinuousModule(ContinuousGenerativeModule):
    """
    Diffusion / flow in a frozen autoencoder's latent space.

    - ``vae`` stage 1: the encoder's posterior mean is the latent; decode with the decoder.
    - ``vqvae`` stage 1: the quantized embeddings are the latent; sampled latents
      are snapped to the codebook before decoding.

    Dense conditions are average-pooled to the latent length; observation masks
    and training-free inpainting are signal-space features and are not used here
    (a ``masked_signal`` condition still works).
    """

    config_cls = LatentConfig

    def _model_channels(self, config: LatentConfig) -> int:
        if config.autoencoder_family not in ("vae", "vqvae"):
            raise ValueError(
                f"autoencoder_family must be 'vae' or 'vqvae', got {config.autoencoder_family!r}"
            )
        if config.loss_on != "all" or config.clean_observed:
            raise ValueError(
                "loss_on='unobserved' and clean_observed are signal-space options"
            )
        from ecgen.families import VAELightning, VQVAELightning

        cls = VAELightning if config.autoencoder_family == "vae" else VQVAELightning
        if config.autoencoder_config is not None:
            ae = cls(**config.autoencoder_config)
        else:
            if (
                not config.autoencoder_checkpoint
                or not Path(config.autoencoder_checkpoint).exists()
            ):
                raise FileNotFoundError(
                    f"Autoencoder checkpoint not found: {config.autoencoder_checkpoint!r}"
                )
            ae = cls.load_from_checkpoint(
                config.autoencoder_checkpoint, map_location="cpu"
            )
            config.autoencoder_config = dict(ae.config.__dict__)
            self.hparams["autoencoder_config"] = config.autoencoder_config
        if config.autoencoder_family == "vae":
            if ae.config.latent_type != "spatial" or ae.is_conditional:
                raise ValueError(
                    "Latent models need an unconditional VAE with latent_type='spatial'"
                )
            channels = ae.config.latent_channels
        else:
            channels = ae.config.latent_channels
        ae.eval()
        ae.freeze()
        self.autoencoder = ae
        self.downsample = 2 ** (len(ae.config.channel_multipliers) - 1)
        scale = config.latent_scale if config.latent_scale is not None else 0.0
        self.register_buffer("latent_scale", torch.tensor(float(scale)))
        return channels

    def train(self, mode: bool = True):
        super().train(mode)
        self.autoencoder.eval()  # stays frozen (no dropout / batch-stat changes)
        return self

    # ------------------------------------------------------------ latents

    @torch.no_grad()
    def encode(self, x: Tensor) -> Tensor:
        """Unscaled latent of a signal."""
        if self.config.autoencoder_family == "vae":
            mean, _ = self.autoencoder.vae.encode(x)
            return mean
        z, _ = self.autoencoder.vqvae.encoder(x)
        _, z_q, _ = self.autoencoder.vqvae.quantize(z)
        return z_q

    @torch.no_grad()
    def decode(self, z: Tensor) -> Tensor:
        """Signal from an unscaled latent."""
        if self.config.autoencoder_family == "vae":
            return self.autoencoder.vae.decode(z)
        _, z_q, _ = self.autoencoder.vqvae.quantize(z)
        return self.autoencoder.vqvae.decode(z_q)

    def to_model_space(self, x: Tensor) -> Tensor:
        z = self.encode(x)
        if float(self.latent_scale) == 0.0:  # first batch: fix the scale
            self.latent_scale.fill_(1.0 / z.float().std().clamp_min(1e-6))
        return z * self.latent_scale

    def from_model_space(self, z: Tensor) -> Tensor:
        return self.decode(z / self.latent_scale)

    def model_space_shape(self, n: int, seq_length: int) -> tuple[int, ...]:
        if seq_length % self.downsample:
            raise ValueError(f"seq_length must be a multiple of {self.downsample}")
        return (n, self._latent_channels(), seq_length // self.downsample)

    def _latent_channels(self) -> int:
        return self.autoencoder.config.latent_channels

    def configure_optimizers(self) -> Any:
        # the frozen autoencoder has requires_grad=False, so the base filter skips it
        return super().configure_optimizers()

encode

encode(x: Tensor) -> Tensor

Unscaled latent of a signal.

Source code in src/ecgen/families/continuous/latent.py

@torch.no_grad()
def encode(self, x: Tensor) -> Tensor:
    """Unscaled latent of a signal."""
    if self.config.autoencoder_family == "vae":
        mean, _ = self.autoencoder.vae.encode(x)
        return mean
    z, _ = self.autoencoder.vqvae.encoder(x)
    _, z_q, _ = self.autoencoder.vqvae.quantize(z)
    return z_q

decode

decode(z: Tensor) -> Tensor

Signal from an unscaled latent.

Source code in src/ecgen/families/continuous/latent.py

@torch.no_grad()
def decode(self, z: Tensor) -> Tensor:
    """Signal from an unscaled latent."""
    if self.config.autoencoder_family == "vae":
        return self.autoencoder.vae.decode(z)
    _, z_q, _ = self.autoencoder.vqvae.quantize(z)
    return self.autoencoder.vqvae.decode(z_q)

LatentDiffusionModule

Bases: LatentContinuousModule

Latent diffusion (default process ddpm).

Source code in src/ecgen/families/continuous/latent.py

@FAMILIES.register("latent_diffusion")
class LatentDiffusionModule(LatentContinuousModule):
    """Latent diffusion (default process ``ddpm``)."""

    family = "latent_diffusion"

LatentFlowMatchingModule

Bases: LatentContinuousModule

Latent flow matching (default process rectified_flow).

Source code in src/ecgen/families/continuous/latent.py

@FAMILIES.register("latent_flow_matching")
class LatentFlowMatchingModule(LatentContinuousModule):
    """Latent flow matching (default process ``rectified_flow``)."""

    family = "latent_flow_matching"
    config_cls = LatentFlowConfig

ecgen.processes

Corruption processes for continuous generative models (diffusion, flow matching).

A :class:Process defines how training pairs are made and how samples are drawn; the family module (:class:ecgen.families.continuous.ContinuousGenerativeModule) owns the network, conditioning, EMA, and optimisation, so adding a new diffusion/flow method means adding a Process, not a family.

model_fn(x, t) passed to :meth:Process.sample returns the raw network output at time t ([B], in the process's own time units); guidance is applied by the caller (see :func:cfg_model_fn).

TrainingPair dataclass

What one training step needs: the network input, its target, and per-sample weights.

Source code in src/ecgen/processes/base.py

@dataclass
class TrainingPair:
    """What one training step needs: the network input, its target, and per-sample weights."""

    xt: Tensor  # corrupted signal [B, C, T]
    t: Tensor  # network time input [B]
    target: Tensor  # regression target [B, C, T]
    weight: Tensor  # per-sample loss weight [B]

Inpainting dataclass

Training-free inpainting at sampling time.

Attributes:

Name	Type	Description
obs	Tensor	Observed signal [B, C, T].
mask	Tensor	1 = observed.
mode	str	"replace" (overwrite the observed part with a correspondingly noised observation after every step), "repaint" (replace, plus jumps resampling round trips per step; Lugmayr et al., 2022), or "clean" (overwrite with the clean observation; for models trained with clean_observed).
jumps	int	RePaint resampling count per step.

Source code in src/ecgen/processes/base.py

@dataclass
class Inpainting:
    """
    Training-free inpainting at sampling time.

    Attributes:
        obs: Observed signal ``[B, C, T]``.
        mask: 1 = observed.
        mode: ``"replace"`` (overwrite the observed part with a correspondingly
            noised observation after every step), ``"repaint"`` (replace, plus
            ``jumps`` resampling round trips per step; Lugmayr et al., 2022), or
            ``"clean"`` (overwrite with the clean observation; for models trained
            with ``clean_observed``).
        jumps: RePaint resampling count per step.
    """

    obs: Tensor
    mask: Tensor
    mode: str = "replace"
    jumps: int = 5

Process

Bases: Module

Base class; see the module docstring.

Source code in src/ecgen/processes/base.py

class Process(nn.Module):
    """Base class; see the module docstring."""

    #: What the network predicts (``"eps"``, ``"v"``, ``"x0"``, ``"velocity"``).
    prediction: str

    def training_pair(self, x0: Tensor, noise: Tensor) -> TrainingPair:
        raise NotImplementedError

    def sample(
        self,
        model_fn: ModelFn,
        shape: tuple[int, ...],
        steps: int,
        device: torch.device,
        generator: Optional[torch.Generator] = None,
        inpaint: Optional[Inpainting] = None,
        **kwargs,
    ) -> Tensor:
        raise NotImplementedError

cfg_model_fn

cfg_model_fn(model_fn_cond: ModelFn, model_fn_uncond: Optional[ModelFn], guidance: float) -> ModelFn

Classifier-free guidance: u + w (c - u) on the network output.

Source code in src/ecgen/processes/base.py

def cfg_model_fn(
    model_fn_cond: ModelFn, model_fn_uncond: Optional[ModelFn], guidance: float
) -> ModelFn:
    """Classifier-free guidance: ``u + w (c - u)`` on the network output."""
    if model_fn_uncond is None or guidance == 1.0:
        return model_fn_cond

    def fn(x: Tensor, t: Tensor) -> Tensor:
        c = model_fn_cond(x, t)
        u = model_fn_uncond(x, t)
        return u + guidance * (c - u)

    return fn

randn

randn(shape, device, generator: Optional[Generator] = None) -> Tensor

Gaussian noise on device, drawn from generator (on its own device) if given.

Source code in src/ecgen/processes/base.py

def randn(shape, device, generator: Optional[torch.Generator] = None) -> Tensor:
    """Gaussian noise on ``device``, drawn from ``generator`` (on its own device) if given."""
    if generator is None:
        return torch.randn(shape, device=device)
    return torch.randn(shape, generator=generator, device=generator.device).to(device)

Discrete-time Gaussian diffusion (DDPM, Ho et al., 2020) with DDIM sampling.

DDPMProcess

Bases: Process

Parameters:

Name	Type	Description	Default
num_steps	int	Diffusion steps T.	1000
schedule	str	"cosine" or "linear".	'cosine'
prediction	str	Network target, "eps", "v" (Salimans & Ho, 2022) or "x0".	'v'
min_snr_gamma	Optional[float]	Min-SNR loss weighting (Hang et al., 2023); None = uniform.	5.0
clip_x0	Optional[float]	Clamp the predicted x0 to [-clip_x0, clip_x0] while sampling (None = no clamp; z-scored ECGs reach about ±10).	None

Source code in src/ecgen/processes/ddpm.py

@PROCESSES.register("ddpm")
class DDPMProcess(Process):
    """
    Args:
        num_steps: Diffusion steps ``T``.
        schedule: ``"cosine"`` or ``"linear"``.
        prediction: Network target, ``"eps"``, ``"v"`` (Salimans & Ho, 2022) or ``"x0"``.
        min_snr_gamma: Min-SNR loss weighting (Hang et al., 2023); ``None`` = uniform.
        clip_x0: Clamp the predicted ``x0`` to ``[-clip_x0, clip_x0]`` while sampling
            (``None`` = no clamp; z-scored ECGs reach about ±10).
    """

    def __init__(
        self,
        num_steps: int = 1000,
        schedule: str = "cosine",
        prediction: str = "v",
        min_snr_gamma: Optional[float] = 5.0,
        clip_x0: Optional[float] = None,
    ) -> None:
        super().__init__()
        if prediction not in ("eps", "v", "x0"):
            raise ValueError(f"prediction must be eps, v, or x0; got {prediction!r}")
        self.num_steps = num_steps
        self.prediction = prediction
        self.min_snr_gamma = min_snr_gamma
        self.clip_x0 = clip_x0
        self.register_buffer(
            "alphas_cumprod", _schedule(schedule, num_steps).float(), persistent=False
        )

    # ------------------------------------------------------------- helpers

    def _ac(self, t: Tensor) -> Tensor:
        return self.alphas_cumprod[t].view(-1, 1, 1)

    def to_x0_eps(self, out: Tensor, xt: Tensor, t: Tensor) -> tuple[Tensor, Tensor]:
        """Network output at integer steps ``t`` → (predicted x0, predicted noise)."""
        ac = self._ac(t)
        a, s = ac.sqrt(), (1 - ac).sqrt()
        if self.prediction == "eps":
            eps = out
            x0 = (xt - s * eps) / a
        elif self.prediction == "x0":
            x0 = out
            eps = (xt - a * x0) / s
        else:  # v = a * eps - s * x0
            x0 = a * xt - s * out
            eps = s * xt + a * out
        if self.clip_x0 is not None:
            x0 = x0.clamp(-self.clip_x0, self.clip_x0)
            eps = (xt - a * x0) / s
        return x0, eps

    def q_sample(self, x0: Tensor, noise: Tensor, t: Tensor) -> Tensor:
        ac = self._ac(t)
        return ac.sqrt() * x0 + (1 - ac).sqrt() * noise

    # -------------------------------------------------------------- training

    def training_pair(self, x0: Tensor, noise: Tensor) -> TrainingPair:
        b = x0.shape[0]
        t = torch.randint(0, self.num_steps, (b,), device=x0.device)
        xt = self.q_sample(x0, noise, t)
        ac = self._ac(t)
        if self.prediction == "eps":
            target = noise
        elif self.prediction == "x0":
            target = x0
        else:
            target = ac.sqrt() * noise - (1 - ac).sqrt() * x0

        weight = torch.ones(b, device=x0.device)
        if self.min_snr_gamma is not None:
            snr = (ac / (1 - ac)).view(-1)
            clipped = snr.clamp(max=self.min_snr_gamma)
            weight = {"eps": clipped / snr, "v": clipped / (snr + 1), "x0": clipped}[
                self.prediction
            ]
        return TrainingPair(xt=xt, t=t.float(), target=target, weight=weight)

    # -------------------------------------------------------------- sampling

    def timesteps(self, steps: int) -> list[int]:
        """Evenly spaced steps from ``T - 1`` down to 0."""
        steps = max(1, min(steps, self.num_steps))
        return sorted(
            {
                int(round(x))
                for x in torch.linspace(0, self.num_steps - 1, steps).tolist()
            },
            reverse=True,
        )

    @torch.no_grad()
    def sample(
        self,
        model_fn: ModelFn,
        shape: tuple[int, ...],
        steps: int = 50,
        device: torch.device = torch.device("cpu"),
        generator: Optional[torch.Generator] = None,
        inpaint: Optional[Inpainting] = None,
        eta: float = 0.0,
        x_init: Optional[Tensor] = None,
        start_step: Optional[int] = None,
        **kwargs,
    ) -> Tensor:
        """
        DDIM sampling (Song et al., 2021) over ``steps`` evenly spaced steps;
        ``eta=1`` is ancestral (DDPM-like) sampling, ``eta=0`` deterministic.

        ``x_init`` / ``start_step`` start from a partially noised signal (SDEdit).
        """
        ts = self.timesteps(steps)
        if start_step is not None:
            ts = [t for t in ts if t <= start_step]
        x = randn(shape, device, generator) if x_init is None else x_init

        for i, t in enumerate(ts):
            t_prev = ts[i + 1] if i + 1 < len(ts) else -1
            jumps = (
                inpaint.jumps
                if (inpaint is not None and inpaint.mode == "repaint" and t_prev >= 0)
                else 1
            )
            for j in range(jumps):
                x_prev = self._ddim_step(model_fn, x, t, t_prev, eta, device, generator)
                if inpaint is not None:
                    x_prev = self._replace(x_prev, inpaint, t_prev, device, generator)
                if (
                    j < jumps - 1
                ):  # RePaint: diffuse back from t_prev to t and redo the step
                    x = self._renoise(x_prev, t_prev, t, device, generator)
                else:
                    x = x_prev
        return x

    def _ddim_step(self, model_fn, x, t, t_prev, eta, device, generator) -> Tensor:
        tt = torch.full((x.shape[0],), t, device=device, dtype=torch.long)
        x0, eps = self.to_x0_eps(model_fn(x, tt.float()), x, tt)
        if t_prev < 0:
            return x0
        ac, ac_prev = self.alphas_cumprod[t], self.alphas_cumprod[t_prev]
        sigma = (
            eta * ((1 - ac_prev) / (1 - ac) * (1 - ac / ac_prev)).clamp(min=0).sqrt()
        )
        dir_xt = (1 - ac_prev - sigma**2).clamp(min=0).sqrt() * eps
        x_prev = ac_prev.sqrt() * x0 + dir_xt
        if eta > 0:
            x_prev = x_prev + sigma * randn(x.shape, device, generator)
        return x_prev

    def _replace(self, x: Tensor, inp: Inpainting, t: int, device, generator) -> Tensor:
        """Overwrite the observed region with the observation noised to level ``t``."""
        if t < 0 or inp.mode == "clean":
            known = inp.obs
        else:
            tt = torch.full((x.shape[0],), t, device=device, dtype=torch.long)
            known = self.q_sample(inp.obs, randn(x.shape, device, generator), tt)
        m = inp.mask.to(x.dtype)
        return m * known + (1 - m) * x

    def _renoise(self, x: Tensor, t_from: int, t_to: int, device, generator) -> Tensor:
        """Sample ``q(x_{t_to} | x_{t_from})`` for ``t_to > t_from``."""
        ratio = self.alphas_cumprod[t_to] / self.alphas_cumprod[t_from]
        return ratio.sqrt() * x + (1 - ratio).sqrt() * randn(x.shape, device, generator)

to_x0_eps

to_x0_eps(out: Tensor, xt: Tensor, t: Tensor) -> tuple[Tensor, Tensor]

Network output at integer steps t → (predicted x0, predicted noise).

Source code in src/ecgen/processes/ddpm.py

def to_x0_eps(self, out: Tensor, xt: Tensor, t: Tensor) -> tuple[Tensor, Tensor]:
    """Network output at integer steps ``t`` → (predicted x0, predicted noise)."""
    ac = self._ac(t)
    a, s = ac.sqrt(), (1 - ac).sqrt()
    if self.prediction == "eps":
        eps = out
        x0 = (xt - s * eps) / a
    elif self.prediction == "x0":
        x0 = out
        eps = (xt - a * x0) / s
    else:  # v = a * eps - s * x0
        x0 = a * xt - s * out
        eps = s * xt + a * out
    if self.clip_x0 is not None:
        x0 = x0.clamp(-self.clip_x0, self.clip_x0)
        eps = (xt - a * x0) / s
    return x0, eps

timesteps

timesteps(steps: int) -> list[int]

Evenly spaced steps from T - 1 down to 0.

Source code in src/ecgen/processes/ddpm.py

def timesteps(self, steps: int) -> list[int]:
    """Evenly spaced steps from ``T - 1`` down to 0."""
    steps = max(1, min(steps, self.num_steps))
    return sorted(
        {
            int(round(x))
            for x in torch.linspace(0, self.num_steps - 1, steps).tolist()
        },
        reverse=True,
    )

sample

sample(model_fn: ModelFn, shape: tuple[int, ...], steps: int = 50, device: device = torch.device('cpu'), generator: Optional[Generator] = None, inpaint: Optional[Inpainting] = None, eta: float = 0.0, x_init: Optional[Tensor] = None, start_step: Optional[int] = None, **kwargs) -> Tensor

DDIM sampling (Song et al., 2021) over steps evenly spaced steps; eta=1 is ancestral (DDPM-like) sampling, eta=0 deterministic.

x_init / start_step start from a partially noised signal (SDEdit).

Source code in src/ecgen/processes/ddpm.py

@torch.no_grad()
def sample(
    self,
    model_fn: ModelFn,
    shape: tuple[int, ...],
    steps: int = 50,
    device: torch.device = torch.device("cpu"),
    generator: Optional[torch.Generator] = None,
    inpaint: Optional[Inpainting] = None,
    eta: float = 0.0,
    x_init: Optional[Tensor] = None,
    start_step: Optional[int] = None,
    **kwargs,
) -> Tensor:
    """
    DDIM sampling (Song et al., 2021) over ``steps`` evenly spaced steps;
    ``eta=1`` is ancestral (DDPM-like) sampling, ``eta=0`` deterministic.

    ``x_init`` / ``start_step`` start from a partially noised signal (SDEdit).
    """
    ts = self.timesteps(steps)
    if start_step is not None:
        ts = [t for t in ts if t <= start_step]
    x = randn(shape, device, generator) if x_init is None else x_init

    for i, t in enumerate(ts):
        t_prev = ts[i + 1] if i + 1 < len(ts) else -1
        jumps = (
            inpaint.jumps
            if (inpaint is not None and inpaint.mode == "repaint" and t_prev >= 0)
            else 1
        )
        for j in range(jumps):
            x_prev = self._ddim_step(model_fn, x, t, t_prev, eta, device, generator)
            if inpaint is not None:
                x_prev = self._replace(x_prev, inpaint, t_prev, device, generator)
            if (
                j < jumps - 1
            ):  # RePaint: diffuse back from t_prev to t and redo the step
                x = self._renoise(x_prev, t_prev, t, device, generator)
            else:
                x = x_prev
    return x

Flow matching / rectified flow (Lipman et al., 2023; Liu et al., 2023).

Time runs from t = 0 (noise x0 ~ N(0, I)) to t = 1 (data x1):

x_t = t * x1 + (1 - (1 - sigma_min) * t) * x0,   target velocity = x1 - (1 - sigma_min) * x0

and sampling integrates dx/dt = v(x, t) with a fixed-step ODE solver.

RectifiedFlowProcess

Bases: Process

Parameters:

Name	Type	Description	Default
t_sampling	str	"uniform" or "logit_normal" (Esser et al., 2024).	'logit_normal'
logit_mean	float	Mean of the logit-normal time distribution.	0.0
logit_std	float	Standard deviation of the logit-normal time distribution.	1.0
sigma_min	float	Width of the conditional path at t = 1 (0 = rectified flow).	0.0
ot_coupling	bool	Pair noise and data within each batch by optimal transport (minibatch OT-CFM, Tong et al., 2024) for straighter paths.	False
time_scale	float	The network sees t * time_scale (matches the 0..1000 range the sinusoidal embedding was tuned for in DDPM).	1000.0
solver	str	Default ODE solver, "euler" or "heun".	'euler'

Source code in src/ecgen/processes/flow.py

@PROCESSES.register("rectified_flow")
class RectifiedFlowProcess(Process):
    """
    Args:
        t_sampling: ``"uniform"`` or ``"logit_normal"`` (Esser et al., 2024).
        logit_mean: Mean of the logit-normal time distribution.
        logit_std: Standard deviation of the logit-normal time distribution.
        sigma_min: Width of the conditional path at ``t = 1`` (0 = rectified flow).
        ot_coupling: Pair noise and data within each batch by optimal transport
            (minibatch OT-CFM, Tong et al., 2024) for straighter paths.
        time_scale: The network sees ``t * time_scale`` (matches the 0..1000 range
            the sinusoidal embedding was tuned for in DDPM).
        solver: Default ODE solver, ``"euler"`` or ``"heun"``.
    """

    prediction = "velocity"

    def __init__(
        self,
        t_sampling: str = "logit_normal",
        logit_mean: float = 0.0,
        logit_std: float = 1.0,
        sigma_min: float = 0.0,
        ot_coupling: bool = False,
        time_scale: float = 1000.0,
        solver: str = "euler",
    ) -> None:
        super().__init__()
        if t_sampling not in ("uniform", "logit_normal"):
            raise ValueError(
                f"t_sampling must be 'uniform' or 'logit_normal', got {t_sampling!r}"
            )
        if solver not in ("euler", "heun"):
            raise ValueError(f"solver must be 'euler' or 'heun', got {solver!r}")
        self.t_sampling = t_sampling
        self.logit_mean, self.logit_std = logit_mean, logit_std
        self.sigma_min = sigma_min
        self.ot_coupling = ot_coupling
        self.time_scale = time_scale
        self.solver = solver

    # -------------------------------------------------------------- training

    def sample_t(self, b: int, device) -> Tensor:
        if self.t_sampling == "uniform":
            return torch.rand(b, device=device)
        return torch.sigmoid(
            self.logit_mean + self.logit_std * torch.randn(b, device=device)
        )

    def interpolate(self, x1: Tensor, x0: Tensor, t: Tensor) -> Tensor:
        tt = t.view(-1, *([1] * (x1.dim() - 1)))
        return tt * x1 + (1 - (1 - self.sigma_min) * tt) * x0

    def training_pair(self, x1: Tensor, noise: Tensor) -> TrainingPair:
        if self.ot_coupling and x1.shape[0] > 1:
            noise = noise[ot_pairing(x1, noise)]
        t = self.sample_t(x1.shape[0], x1.device)
        return TrainingPair(
            xt=self.interpolate(x1, noise, t),
            t=t * self.time_scale,
            target=x1 - (1 - self.sigma_min) * noise,
            weight=torch.ones_like(t),
        )

    # -------------------------------------------------------------- sampling

    @torch.no_grad()
    def sample(
        self,
        model_fn: ModelFn,
        shape: tuple[int, ...],
        steps: int = 50,
        device: torch.device = torch.device("cpu"),
        generator: Optional[torch.Generator] = None,
        inpaint: Optional[Inpainting] = None,
        solver: Optional[str] = None,
        x_init: Optional[Tensor] = None,
        t_start: float = 0.0,
        **kwargs,
    ) -> Tensor:
        """
        Integrate from noise (``t = t_start``) to data (``t = 1``) in ``steps`` steps.

        Inpainting: ``"replace"`` overwrites the observed part with the point on
        the straight path from the initial noise to the observation; ``"clean"``
        with the observation itself. (RePaint is a diffusion method; use ``replace``.)
        ``x_init`` / ``t_start`` start mid-path (SDEdit-style editing).
        """
        solver = solver or self.solver
        if inpaint is not None and inpaint.mode == "repaint":
            raise ValueError("RePaint is not defined for flows; use inpaint='replace'")
        noise = randn(shape, device, generator)
        x = noise if x_init is None else x_init
        ts = torch.linspace(t_start, 1.0, steps + 1, device=device)

        def v(x_: Tensor, t: Tensor) -> Tensor:
            return model_fn(
                x_, torch.full((shape[0],), float(t) * self.time_scale, device=device)
            )

        def known(t: Tensor) -> Optional[Tensor]:
            if inpaint is None:
                return None
            if inpaint.mode == "clean":
                return inpaint.obs
            return self.interpolate(inpaint.obs, noise, t.expand(shape[0]))

        for i in range(steps):
            t, t_next = ts[i], ts[i + 1]
            h = t_next - t
            if inpaint is not None:
                m = inpaint.mask.to(x.dtype)
                x = m * known(t) + (1 - m) * x
            d = v(x, t)
            if solver == "heun" and i < steps - 1:
                x_euler = x + h * d
                x = x + h * 0.5 * (d + v(x_euler, t_next))
            else:
                x = x + h * d
        if inpaint is not None:
            m = inpaint.mask.to(x.dtype)
            x = m * inpaint.obs + (1 - m) * x
        return x

sample

sample(model_fn: ModelFn, shape: tuple[int, ...], steps: int = 50, device: device = torch.device('cpu'), generator: Optional[Generator] = None, inpaint: Optional[Inpainting] = None, solver: Optional[str] = None, x_init: Optional[Tensor] = None, t_start: float = 0.0, **kwargs) -> Tensor

Integrate from noise (t = t_start) to data (t = 1) in steps steps.

Inpainting: "replace" overwrites the observed part with the point on the straight path from the initial noise to the observation; "clean" with the observation itself. (RePaint is a diffusion method; use replace.) x_init / t_start start mid-path (SDEdit-style editing).

Source code in src/ecgen/processes/flow.py

@torch.no_grad()
def sample(
    self,
    model_fn: ModelFn,
    shape: tuple[int, ...],
    steps: int = 50,
    device: torch.device = torch.device("cpu"),
    generator: Optional[torch.Generator] = None,
    inpaint: Optional[Inpainting] = None,
    solver: Optional[str] = None,
    x_init: Optional[Tensor] = None,
    t_start: float = 0.0,
    **kwargs,
) -> Tensor:
    """
    Integrate from noise (``t = t_start``) to data (``t = 1``) in ``steps`` steps.

    Inpainting: ``"replace"`` overwrites the observed part with the point on
    the straight path from the initial noise to the observation; ``"clean"``
    with the observation itself. (RePaint is a diffusion method; use ``replace``.)
    ``x_init`` / ``t_start`` start mid-path (SDEdit-style editing).
    """
    solver = solver or self.solver
    if inpaint is not None and inpaint.mode == "repaint":
        raise ValueError("RePaint is not defined for flows; use inpaint='replace'")
    noise = randn(shape, device, generator)
    x = noise if x_init is None else x_init
    ts = torch.linspace(t_start, 1.0, steps + 1, device=device)

    def v(x_: Tensor, t: Tensor) -> Tensor:
        return model_fn(
            x_, torch.full((shape[0],), float(t) * self.time_scale, device=device)
        )

    def known(t: Tensor) -> Optional[Tensor]:
        if inpaint is None:
            return None
        if inpaint.mode == "clean":
            return inpaint.obs
        return self.interpolate(inpaint.obs, noise, t.expand(shape[0]))

    for i in range(steps):
        t, t_next = ts[i], ts[i + 1]
        h = t_next - t
        if inpaint is not None:
            m = inpaint.mask.to(x.dtype)
            x = m * known(t) + (1 - m) * x
        d = v(x, t)
        if solver == "heun" and i < steps - 1:
            x_euler = x + h * d
            x = x + h * 0.5 * (d + v(x_euler, t_next))
        else:
            x = x + h * d
    if inpaint is not None:
        m = inpaint.mask.to(x.dtype)
        x = m * inpaint.obs + (1 - m) * x
    return x

ot_pairing

ot_pairing(x1: Tensor, x0: Tensor) -> Tensor

Permutation of x0 minimising the summed squared distance to x1 (exact assignment).

Source code in src/ecgen/processes/flow.py

def ot_pairing(x1: Tensor, x0: Tensor) -> Tensor:
    """Permutation of ``x0`` minimising the summed squared distance to ``x1`` (exact assignment)."""
    from scipy.optimize import linear_sum_assignment

    cost = torch.cdist(x1.flatten(1).float(), x0.flatten(1).float()) ** 2
    _, cols = linear_sum_assignment(cost.detach().cpu().numpy())
    return torch.as_tensor(cols, device=x0.device)

ecgen.nn.backbones

1D U-Net denoiser / velocity network (ADM-style) for ECGs.

forward(x, t, cond) takes the noisy signal [B, C, T], the time [B] (any float scale; it is embedded sinusoidally) and an optional :class:~ecgen.types.ConditionOutput:

• cond.vector is added to the time embedding (AddToTimeEmb), which modulates every residual block through AdaGN;
• cond.tokens are attended to by cross-attention at the attention levels;
• cond.dense is concatenated to the input.

ResBlock

Bases: Module

GN-SiLU-Conv, AdaGN(time+condition)-SiLU-Dropout-Conv(zero init), residual.

Source code in src/ecgen/nn/backbones/unet1d.py

class ResBlock(nn.Module):
    """GN-SiLU-Conv, AdaGN(time+condition)-SiLU-Dropout-Conv(zero init), residual."""

    def __init__(
        self,
        in_ch: int,
        out_ch: int,
        emb_dim: int,
        dropout: float = 0.0,
        groups: int = 8,
    ):
        super().__init__()
        self.norm1 = nn.GroupNorm(groups, in_ch)
        self.conv1 = nn.Conv1d(in_ch, out_ch, 3, padding=1)
        self.norm2 = AdaGN(emb_dim, out_ch, groups)
        self.drop = nn.Dropout(dropout)
        self.conv2 = _zero(nn.Conv1d(out_ch, out_ch, 3, padding=1))
        self.skip = nn.Conv1d(in_ch, out_ch, 1) if in_ch != out_ch else nn.Identity()

    def forward(self, x: Tensor, emb: Tensor) -> Tensor:
        h = self.conv1(F.silu(self.norm1(x)))
        h = self.conv2(self.drop(F.silu(self.norm2(h, emb))))
        return self.skip(x) + h

UNet1d

Bases: Module

Parameters:

Name	Type	Description	Default
in_channels	int	Signal channels (leads).	12
out_channels	Optional[int]	Output channels (default in_channels).	None
base_channels	int	Width of the first level.	64
channel_mults	Sequence[int]	Per-level width multipliers; the signal length must be divisible by 2 ** (len(channel_mults) - 1).	(1, 2, 4, 4)
num_res_blocks	int	Residual blocks per level.	2
attn_levels	Sequence[int]	Levels (0 = full resolution) with self-attention (and cross-attention when context_dim > 0).	(3,)
heads	int	Attention heads.	4
dropout	float	Dropout inside residual blocks.	0.0
cond_dim	int	Size of ConditionOutput.vector (0 = none).	0
context_dim	int	Size of ConditionOutput.tokens (0 = none).	0
dense_channels	int	Channels of ConditionOutput.dense (0 = none).	0

Source code in src/ecgen/nn/backbones/unet1d.py

@BACKBONES.register("unet1d")
class UNet1d(nn.Module):
    """
    Args:
        in_channels: Signal channels (leads).
        out_channels: Output channels (default ``in_channels``).
        base_channels: Width of the first level.
        channel_mults: Per-level width multipliers; the signal length must be
            divisible by ``2 ** (len(channel_mults) - 1)``.
        num_res_blocks: Residual blocks per level.
        attn_levels: Levels (0 = full resolution) with self-attention (and
            cross-attention when ``context_dim`` > 0).
        heads: Attention heads.
        dropout: Dropout inside residual blocks.
        cond_dim: Size of ``ConditionOutput.vector`` (0 = none).
        context_dim: Size of ``ConditionOutput.tokens`` (0 = none).
        dense_channels: Channels of ``ConditionOutput.dense`` (0 = none).
    """

    accepts = frozenset({"vector", "tokens", "dense"})

    def __init__(
        self,
        in_channels: int = 12,
        out_channels: Optional[int] = None,
        base_channels: int = 64,
        channel_mults: Sequence[int] = (1, 2, 4, 4),
        num_res_blocks: int = 2,
        attn_levels: Sequence[int] = (3,),
        heads: int = 4,
        dropout: float = 0.0,
        cond_dim: int = 0,
        context_dim: int = 0,
        dense_channels: int = 0,
    ) -> None:
        super().__init__()
        out_channels = out_channels or in_channels
        self.levels = len(channel_mults)
        self.downsample_factor = 2 ** (self.levels - 1)
        self.time_embed = TimestepEmbedding(base_channels)
        emb_dim = self.time_embed.out_dim
        self.cond_embed = AddToTimeEmb(cond_dim, emb_dim) if cond_dim else None
        self.conv_in = nn.Conv1d(
            in_channels + dense_channels, base_channels, 3, padding=1
        )

        widths = [base_channels * m for m in channel_mults]
        self.down = nn.ModuleList()
        self.downsample = nn.ModuleList()
        skips = [base_channels]
        ch = base_channels
        for lvl, w in enumerate(widths):
            chans = []
            for _ in range(num_res_blocks):
                chans.append((ch, w))
                ch = w
                skips.append(ch)
            self.down.append(
                _Level(chans, emb_dim, dropout, lvl in attn_levels, heads, context_dim)
            )
            if lvl < self.levels - 1:
                self.downsample.append(nn.Conv1d(ch, ch, 3, stride=2, padding=1))
                skips.append(ch)

        self.mid = _Level(
            [(ch, ch), (ch, ch)], emb_dim, dropout, True, heads, context_dim
        )

        self.up = nn.ModuleList()
        self.upsample = nn.ModuleList()
        for lvl, w in reversed(list(enumerate(widths))):
            chans = []
            for _ in range(num_res_blocks + 1):
                chans.append((ch + skips.pop(), w))
                ch = w
            self.up.append(
                _Level(chans, emb_dim, dropout, lvl in attn_levels, heads, context_dim)
            )
            if lvl > 0:
                self.upsample.append(nn.Conv1d(ch, ch, 3, padding=1))

        self.norm_out = nn.GroupNorm(8, ch)
        self.conv_out = _zero(nn.Conv1d(ch, out_channels, 3, padding=1))

    @staticmethod
    def _run(level: _Level, i: int, h: Tensor, emb: Tensor, tokens, mask) -> Tensor:
        h = level.blocks[i](h, emb)
        h = level.attn[i](h)
        if isinstance(level.cross[i], CrossAttention1d):
            h = level.cross[i](h, tokens, mask)
        return h

    def forward(
        self, x: Tensor, t: Tensor, cond: Optional[ConditionOutput] = None
    ) -> Tensor:
        if x.shape[-1] % self.downsample_factor:
            raise ValueError(
                f"Signal length {x.shape[-1]} must be divisible by {self.downsample_factor} "
                f"(2 ** (len(channel_mults) - 1))"
            )
        emb = self.time_embed(t)
        tokens = mask = None
        if cond is not None:
            if self.cond_embed is not None:
                emb = self.cond_embed(emb, cond.vector)
            tokens, mask = cond.tokens, cond.tokens_mask
            x = concat_dense(x, cond.dense)

        h = self.conv_in(x)
        hs = [h]
        for lvl, level in enumerate(self.down):
            for i in range(len(level.blocks)):
                h = self._run(level, i, h, emb, tokens, mask)
                hs.append(h)
            if lvl < self.levels - 1:
                h = self.downsample[lvl](h)
                hs.append(h)

        for i in range(len(self.mid.blocks)):
            h = self._run(self.mid, i, h, emb, tokens, mask)

        for j, level in enumerate(self.up):
            for i in range(len(level.blocks)):
                h = self._run(
                    level, i, torch.cat([h, hs.pop()], dim=1), emb, tokens, mask
                )
            if j < self.levels - 1:
                h = self.upsample[j](F.interpolate(h, scale_factor=2.0, mode="nearest"))

        return self.conv_out(F.silu(self.norm_out(h)))

DiT-style transformer for 1D signals (Peebles & Xie, 2023).

The signal (plus dense conditions) is cut into non-overlapping patches of patch_size samples, embedded as tokens with fixed sinusoidal positions, and processed by transformer blocks whose LayerNorms are modulated by the time + condition embedding (AdaLN-Zero); condition tokens enter through cross-attention in every block.

DiT1d

Bases: Module

Parameters:

Name	Type	Description	Default
in_channels	int	Signal (or latent) channels.	12
out_channels	Optional[int]	Default in_channels.	None
patch_size	int	Samples per token; the length must be divisible by it.	25
dim	int	Token width.	256
depth	int	Transformer blocks.	8
heads	int	Attention heads.	8
mlp_ratio	float	MLP width / dim.	4.0
cond_dim	int	Size of the condition vector (0 = none).	0
context_dim	int	Size of condition tokens (0 = none).	0
dense_channels	int	Channels of dense conditions (0 = none).	0

Source code in src/ecgen/nn/backbones/dit1d.py

@BACKBONES.register("dit1d")
class DiT1d(nn.Module):
    """
    Args:
        in_channels: Signal (or latent) channels.
        out_channels: Default ``in_channels``.
        patch_size: Samples per token; the length must be divisible by it.
        dim: Token width.
        depth: Transformer blocks.
        heads: Attention heads.
        mlp_ratio: MLP width / ``dim``.
        cond_dim: Size of the condition vector (0 = none).
        context_dim: Size of condition tokens (0 = none).
        dense_channels: Channels of dense conditions (0 = none).
    """

    accepts = frozenset({"vector", "tokens", "dense"})

    def __init__(
        self,
        in_channels: int = 12,
        out_channels: Optional[int] = None,
        patch_size: int = 25,
        dim: int = 256,
        depth: int = 8,
        heads: int = 8,
        mlp_ratio: float = 4.0,
        cond_dim: int = 0,
        context_dim: int = 0,
        dense_channels: int = 0,
    ) -> None:
        super().__init__()
        out_channels = out_channels or in_channels
        self.patch_size = patch_size
        self.out_channels = out_channels
        self.dim = dim
        self.patch = nn.Conv1d(
            in_channels + dense_channels, dim, patch_size, stride=patch_size
        )
        self.time_embed = TimestepEmbedding(dim)
        emb_dim = self.time_embed.out_dim
        self.cond_proj = nn.Linear(cond_dim, emb_dim) if cond_dim else None
        if self.cond_proj is not None:
            nn.init.zeros_(self.cond_proj.weight)
            nn.init.zeros_(self.cond_proj.bias)
        self.blocks = nn.ModuleList(
            DiTBlock(dim, heads, emb_dim, context_dim, mlp_ratio) for _ in range(depth)
        )
        self.final = AdaLNZero(emb_dim, dim)
        self.out = nn.Linear(dim, patch_size * out_channels)
        nn.init.zeros_(self.out.weight)
        nn.init.zeros_(self.out.bias)

    def forward(
        self, x: Tensor, t: Tensor, cond: Optional[ConditionOutput] = None
    ) -> Tensor:
        b, _, length = x.shape
        if length % self.patch_size:
            raise ValueError(
                f"Signal length {length} must be divisible by patch_size={self.patch_size}"
            )
        emb = self.time_embed(t)
        tokens = mask = None
        if cond is not None:
            if self.cond_proj is not None and cond.vector is not None:
                emb = emb + self.cond_proj(cond.vector)
            tokens, mask = cond.tokens, cond.tokens_mask
            x = concat_dense(x, cond.dense)
        emb = F.silu(emb)

        h = self.patch(x).transpose(1, 2)  # [B, N, dim]
        h = h + _positions(h.shape[1], self.dim, h.device)
        for block in self.blocks:
            h = block(h, emb, tokens, mask)
        h, _ = self.final(h, emb)
        out = self.out(h)  # [B, N, patch * C]
        out = out.view(b, -1, self.patch_size, self.out_channels).permute(0, 3, 1, 2)
        return out.reshape(b, self.out_channels, length)

State-space backbone (SSSD-style; Alcaraz & Strodthoff, 2023) built from S4D layers (Gu et al., 2022) in pure PyTorch.

An S4D layer is a diagonal linear state-space model applied per channel as a long convolution: its kernel K[l] = 2 Re(sum_n C_n (exp(dt A_n) - 1) / A_n * exp(dt A_n l)) is computed once per forward pass and applied with an FFT, so the cost is O(L log L) and the receptive field is the whole signal. bidirectional adds a second kernel run over the time-reversed input (a non-causal model, as denoisers should be).

S4D

Bases: Module

Diagonal state-space layer [B, H, L] -> [B, H, L].

Parameters:

Name	Type	Description	Default
channels	int	H.	required
state_size	int	N; N / 2 complex conjugate pairs are kept.	64
dt_min	float	Smallest initial step size (learned, log-uniform initialised).	0.001
dt_max	float	Largest initial step size.	0.1
bidirectional	bool	Also run a kernel over the reversed sequence.	True

Source code in src/ecgen/nn/backbones/ssm.py

class S4D(nn.Module):
    """
    Diagonal state-space layer ``[B, H, L] -> [B, H, L]``.

    Args:
        channels: ``H``.
        state_size: ``N``; ``N / 2`` complex conjugate pairs are kept.
        dt_min: Smallest initial step size (learned, log-uniform initialised).
        dt_max: Largest initial step size.
        bidirectional: Also run a kernel over the reversed sequence.
    """

    def __init__(
        self,
        channels: int,
        state_size: int = 64,
        dt_min: float = 1e-3,
        dt_max: float = 1e-1,
        bidirectional: bool = True,
    ) -> None:
        super().__init__()
        n = state_size // 2
        k = 2 if bidirectional else 1
        self.bidirectional = bidirectional
        log_dt = torch.rand(k, channels) * (
            math.log(dt_max) - math.log(dt_min)
        ) + math.log(dt_min)
        self.log_dt = nn.Parameter(log_dt)
        # S4D-Lin initialisation: A_n = -1/2 + i * pi * n
        self.log_a_real = nn.Parameter(torch.log(0.5 * torch.ones(k, channels, n)))
        self.a_imag = nn.Parameter(
            math.pi * torch.arange(n).float().expand(k, channels, n).clone()
        )
        self.c = nn.Parameter(
            torch.view_as_real(
                torch.randn(k, channels, n, dtype=torch.cfloat) * (0.5**0.5)
            )
        )
        self.d = nn.Parameter(torch.randn(channels))

    def kernel(self, length: int) -> Tensor:
        """``[k, H, L]`` convolution kernels."""
        dt = self.log_dt.exp()[..., None]  # [k, H, 1]
        a = -self.log_a_real.exp() + 1j * self.a_imag  # [k, H, N]
        c = torch.view_as_complex(self.c)
        dta = dt * a
        coef = c * (torch.exp(dta) - 1.0) / a  # [k, H, N]
        vander = torch.exp(
            dta[..., None] * torch.arange(length, device=a.device)
        )  # [k, H, N, L]
        return 2 * torch.einsum("khn,khnl->khl", coef, vander).real

    def forward(self, u: Tensor) -> Tensor:
        length = u.shape[-1]
        k = self.kernel(length).to(u.dtype)
        n_fft = 2 * length
        u_f = torch.fft.rfft(u.float(), n=n_fft)
        y = torch.fft.irfft(u_f * torch.fft.rfft(k[0].float(), n=n_fft), n=n_fft)[
            ..., :length
        ]
        if self.bidirectional:
            u_rev = torch.fft.rfft(u.flip(-1).float(), n=n_fft)
            y_rev = torch.fft.irfft(
                u_rev * torch.fft.rfft(k[1].float(), n=n_fft), n=n_fft
            )[..., :length]
            y = y + y_rev.flip(-1)
        return (y + u.float() * self.d[:, None]).to(u.dtype)

kernel

kernel(length: int) -> Tensor

[k, H, L] convolution kernels.

Source code in src/ecgen/nn/backbones/ssm.py

def kernel(self, length: int) -> Tensor:
    """``[k, H, L]`` convolution kernels."""
    dt = self.log_dt.exp()[..., None]  # [k, H, 1]
    a = -self.log_a_real.exp() + 1j * self.a_imag  # [k, H, N]
    c = torch.view_as_complex(self.c)
    dta = dt * a
    coef = c * (torch.exp(dta) - 1.0) / a  # [k, H, N]
    vander = torch.exp(
        dta[..., None] * torch.arange(length, device=a.device)
    )  # [k, H, N, L]
    return 2 * torch.einsum("khn,khnl->khl", coef, vander).real

SSMResidualBlock

Bases: Module

SSSD residual block: (+ time/condition) -> S4D -> gated activation -> residual + skip.

Source code in src/ecgen/nn/backbones/ssm.py

class SSMResidualBlock(nn.Module):
    """SSSD residual block: (+ time/condition) -> S4D -> gated activation -> residual + skip."""

    def __init__(
        self,
        channels: int,
        emb_dim: int,
        state_size: int,
        bidirectional: bool,
        dropout: float,
    ):
        super().__init__()
        self.emb = nn.Linear(emb_dim, channels)
        self.norm = nn.GroupNorm(8, channels)
        self.ssm = S4D(channels, state_size, bidirectional=bidirectional)
        self.mix = nn.Conv1d(channels, 2 * channels, 1)
        self.drop = nn.Dropout(dropout)
        self.res = nn.Conv1d(channels, channels, 1)
        self.skip = nn.Conv1d(channels, channels, 1)

    def forward(self, x: Tensor, emb: Tensor) -> tuple[Tensor, Tensor]:
        h = self.norm(x + self.emb(emb)[..., None])
        h = self.mix(self.ssm(h))
        a, g = h.chunk(2, dim=1)
        h = self.drop(torch.tanh(a) * torch.sigmoid(g))
        return (x + self.res(h)) / math.sqrt(2.0), self.skip(h)

SSMBackbone

Bases: Module

Stack of S4D residual blocks at full resolution (no down-sampling).

Parameters:

Name	Type	Description	Default
in_channels	int	Signal (or latent) channels.	12
out_channels	Optional[int]	Default in_channels.	None
channels	int	Width.	128
depth	int	Residual blocks.	8
state_size	int	S4D state size N.	64
bidirectional	bool	Non-causal (two-sided) kernels.	True
dropout	float	Dropout after the gate.	0.0
cond_dim	int	Size of the condition vector, added to the time embedding (0 = none).	0
dense_channels	int	Channels of dense conditions, concatenated to the input. Tokens are not supported.	0

Source code in src/ecgen/nn/backbones/ssm.py

@BACKBONES.register("ssm")
class SSMBackbone(nn.Module):
    """
    Stack of S4D residual blocks at full resolution (no down-sampling).

    Args:
        in_channels: Signal (or latent) channels.
        out_channels: Default ``in_channels``.
        channels: Width.
        depth: Residual blocks.
        state_size: S4D state size ``N``.
        bidirectional: Non-causal (two-sided) kernels.
        dropout: Dropout after the gate.
        cond_dim: Size of the condition vector, added to the time embedding (0 = none).
        dense_channels: Channels of dense conditions, concatenated to the input.
            Tokens are not supported.
    """

    accepts = frozenset({"vector", "dense"})

    def __init__(
        self,
        in_channels: int = 12,
        out_channels: Optional[int] = None,
        channels: int = 128,
        depth: int = 8,
        state_size: int = 64,
        bidirectional: bool = True,
        dropout: float = 0.0,
        cond_dim: int = 0,
        context_dim: int = 0,
        dense_channels: int = 0,
    ) -> None:
        super().__init__()
        if context_dim:
            raise ValueError(
                "The ssm backbone does not take token conditions; use unet1d or dit1d"
            )
        out_channels = out_channels or in_channels
        self.time_embed = TimestepEmbedding(channels)
        emb_dim = self.time_embed.out_dim
        self.cond_embed = AddToTimeEmb(cond_dim, emb_dim) if cond_dim else None
        self.conv_in = nn.Conv1d(in_channels + dense_channels, channels, 1)
        self.blocks = nn.ModuleList(
            SSMResidualBlock(channels, emb_dim, state_size, bidirectional, dropout)
            for _ in range(depth)
        )
        self.conv_out = nn.Sequential(
            nn.SiLU(),
            nn.Conv1d(channels, channels, 1),
            nn.SiLU(),
            nn.Conv1d(channels, out_channels, 1),
        )
        nn.init.zeros_(self.conv_out[-1].weight)
        nn.init.zeros_(self.conv_out[-1].bias)

    def forward(
        self, x: Tensor, t: Tensor, cond: Optional[ConditionOutput] = None
    ) -> Tensor:
        emb = self.time_embed(t)
        if cond is not None:
            if self.cond_embed is not None:
                emb = self.cond_embed(emb, cond.vector)
            x = concat_dense(x, cond.dense)
        emb = F.silu(emb)
        h = self.conv_in(x)
        skip = torch.zeros_like(h)
        for block in self.blocks:
            h, s = block(h, emb)
            skip = skip + s
        return self.conv_out(skip / math.sqrt(len(self.blocks)))

MambaBackbone

Bases: SSMBackbone

:class:SSMBackbone with bidirectional Mamba layers (Gu & Dao, 2023) in place of S4D. Needs the mamba-ssm package (CUDA only; ecgen[ssm]).

Source code in src/ecgen/nn/backbones/ssm.py

@BACKBONES.register("mamba")
class MambaBackbone(SSMBackbone):
    """
    :class:`SSMBackbone` with bidirectional Mamba layers (Gu & Dao, 2023) in
    place of S4D. Needs the ``mamba-ssm`` package (CUDA only; ``ecgen[ssm]``).
    """

    def __init__(self, *args, state_size: int = 16, **kwargs) -> None:
        try:
            import mamba_ssm  # noqa: F401
        except ImportError as e:  # pragma: no cover - depends on the environment
            raise ImportError(
                "The mamba backbone needs `pip install mamba-ssm` (CUDA)"
            ) from e
        super().__init__(*args, state_size=state_size, **kwargs)
        channels = self.conv_in.out_channels
        for block in self.blocks:
            block.ssm = _BiMamba(channels, state_size)

ecgen.training.ema

Exponential moving average of model weights.

EMA

Keeps shadow = decay * shadow + (1 - decay) * param for the parameters of some modules; :meth:swap temporarily loads the averaged weights.

The shadow is a plain dict of tensors (not a submodule), so the owner stores it in checkpoints itself (:meth:state_dict / :meth:load_state_dict).

Source code in src/ecgen/training/ema.py

class EMA:
    """
    Keeps ``shadow = decay * shadow + (1 - decay) * param`` for the parameters
    of some modules; :meth:`swap` temporarily loads the averaged weights.

    The shadow is a plain dict of tensors (not a submodule), so the owner
    stores it in checkpoints itself (:meth:`state_dict` / :meth:`load_state_dict`).
    """

    def __init__(self, modules: Sequence[nn.Module], decay: float = 0.999) -> None:
        self.decay = decay
        self.modules = [m for m in modules if m is not None]
        self.shadow: Dict[str, Tensor] = {
            name: p.detach().clone() for name, p in self._named_params()
        }

    def _named_params(self) -> Iterator[tuple[str, nn.Parameter]]:
        for i, m in enumerate(self.modules):
            for name, p in m.named_parameters():
                if p.requires_grad:
                    yield f"{i}.{name}", p

    @torch.no_grad()
    def update(self) -> None:
        for name, p in self._named_params():
            s = self.shadow[name]
            if s.device != p.device:
                s = self.shadow[name] = s.to(p.device)
            s.lerp_(p.detach(), 1.0 - self.decay)

    @contextmanager
    def swap(self) -> Iterator[None]:
        """Use the averaged weights inside the ``with`` block."""
        backup = {}
        with torch.no_grad():
            for name, p in self._named_params():
                backup[name] = p.detach().clone()
                p.copy_(self.shadow[name].to(p.device))
        try:
            yield
        finally:
            with torch.no_grad():
                for name, p in self._named_params():
                    p.copy_(backup[name])

    def state_dict(self) -> Dict[str, Tensor]:
        return {k: v.detach().cpu() for k, v in self.shadow.items()}

    def load_state_dict(self, state: Dict[str, Tensor]) -> None:
        missing = set(self.shadow) - set(state)
        if missing:
            raise KeyError(f"EMA state is missing {sorted(missing)[:5]}...")
        self.shadow = {k: state[k].clone() for k in self.shadow}

swap

swap() -> Iterator[None]

Use the averaged weights inside the with block.

Source code in src/ecgen/training/ema.py

@contextmanager
def swap(self) -> Iterator[None]:
    """Use the averaged weights inside the ``with`` block."""
    backup = {}
    with torch.no_grad():
        for name, p in self._named_params():
            backup[name] = p.detach().clone()
            p.copy_(self.shadow[name].to(p.device))
    try:
        yield
    finally:
        with torch.no_grad():
            for name, p in self._named_params():
                p.copy_(backup[name])