<filename>pytorch_generative/models/vae/vaes.py
"""Common modules used by Variational Autoencoders."""

import numpy as np
import torch
from torch import nn

from pytorch_generative import nn as pg_nn


@torch.jit.script
def to_var(log_std):
    """Returns the variance given the log standard deviation."""
    return log_std.exp().pow(2)


@torch.jit.script
def unit_gaussian_kl_div(mean, log_std):
    """Returns `KL(p || N(0, 1))` where `p` is a Gaussian with diagonal covariance."""
    return -0.5 * (1 + 2 * log_std - to_var(log_std) - mean ** 2)


@torch.jit.script
def gaussian_kl_div(p_mean, p_log_std, q_mean, q_log_std):
    """Returns `KL(p || q)` where `p` and `q` are Gaussians with diagonal covariance."""
    mean_delta, log_std_delta = (p_mean - q_mean) ** 2, q_log_std - p_log_std
    p_var, q_var = to_var(p_log_std), 2 * to_var(q_log_std)
    return -0.5 + log_std_delta + (p_var + mean_delta) / q_var


@torch.jit.script
def sample_from_gaussian(mu, log_sig):
    """Returns a sample from a Gaussian with diagonal covariance."""
    return mu + log_sig.exp() * torch.randn_like(log_sig)


@torch.jit.script
def _unflatten_tril(x):
    """Unflattens a vector into a lower triangular matrix of shape `dim x dim`."""
    n, dim = x.shape
    idxs = torch.tril_indices(dim, dim)
    tril = torch.zeros(n, dim, dim)
    tril[:, idxs[0, :], idxs[1, :]] = x
    return tril


@torch.jit.script
def gaussian_log_prob(x, mu, chol_sig):
    """Returns the log likelihood of `x` under a Gaussian with full covariance.

    The covariance is assumed to be positive-definite and to have been decomposed into
    the Cholesky factor `chol_sig`, i.e. `sig = chol_sig @ chol_sig.T`.

    Args:
        x: The input for which to compute the log likelihood.
        mu: The mean of the multivariate Gaussian.
        chol_sig: The flattened lower-triangular Cholesky factor of the covariance.
    Returns:
        The log likelihood.
    """
    dim = x.shape[0]
    chol_sig = _unflatten_tril(chol_sig)
    sig = chol_sig @ chol_sig.T
    const = -0.5 * dim * torch.log(torch.tensor(2 * np.pi))
    log_det = -0.5 * torch.logdet(x)
    exp = -0.5 * ((x - mu).T @ sig.inverse() @ (x - mu))
    return const + log_det + exp


class ResidualBlock(nn.Module):
    """A simple residual block."""

    def __init__(self, n_channels, hidden_channels):
        """Initializes a new ResidualBlock instance.

        Args:
            n_channels: Number of input and output channels.
            hidden_channels: Number of hidden channels.
        """
        super().__init__()
        self._net = nn.Sequential(
            nn.ReLU(),
            nn.Conv2d(
                in_channels=n_channels,
                out_channels=hidden_channels,
                kernel_size=3,
                padding=1,
            ),
            nn.ReLU(),
            nn.Conv2d(
                in_channels=hidden_channels, out_channels=n_channels, kernel_size=1
            ),
        )

    def forward(self, x):
        return x + self._net(x)


class ResidualStack(nn.Module):
    """A stack of multiple ResidualBlocks."""

    def __init__(self, n_channels, hidden_channels, n_residual_blocks=1):
        """Initializes a new ResidualStack instance.

        Args:
            n_channels: Number of input and output channels.
            hidden_channels: Number of hidden channels.
            n_residual_blocks: Number of residual blocks in the stack.
        """
        super().__init__()
        self._net = nn.Sequential(
            *[
                ResidualBlock(n_channels, hidden_channels)
                for _ in range(n_residual_blocks)
            ]
            + [nn.ReLU()]
        )

    def forward(self, x):
        return self._net(x)


class Encoder(nn.Module):
    """A feedforward encoder which downsamples its input."""

    def __init__(
        self,
        in_channels,
        out_channels,
        hidden_channels,
        n_residual_blocks,
        residual_channels,
        stride,
    ):
        """Initializes a new Encoder instance.

        Args:
            in_channels: Number of input channels.
            out_channels: Number of output channels.
            hidden_channels: Number of channels in non residual block hidden layers.
            n_residual_blocks: Number of residual blocks in each residual stack.
            residual_channels: Number of hidden channels in residual blocks.
            stride: Stride to use in the downsampling convolutions. Must be even.
        """
        super().__init__()
        assert stride % 2 == 0, '"stride" must be even.'

        net = []
        for i in range(stride // 2):
            first, last = 0, stride // 2 - 1
            in_c = in_channels if i == first else hidden_channels // 2
            out_c = hidden_channels // 2 if i < last else hidden_channels
            net.append(
                nn.Conv2d(
                    in_channels=in_c,
                    out_channels=out_c,
                    kernel_size=4,
                    stride=2,
                    padding=1,
                )
            )
            net.append(nn.ReLU())
        net.append(
            ResidualStack(
                n_channels=hidden_channels,
                hidden_channels=residual_channels,
                n_residual_blocks=n_residual_blocks,
            )
        )
        net.append(
            nn.Conv2d(
                in_channels=hidden_channels,
                out_channels=out_channels,
                kernel_size=3,
                padding=1,
            )
        )
        self._net = nn.Sequential(*net)

    def forward(self, x):
        return self._net(x)


class Decoder(nn.Module):
    """A feedforward encoder which upsamples its input."""

    def __init__(
        self,
        in_channels,
        out_channels,
        hidden_channels,
        n_residual_blocks,
        residual_channels,
        stride,
    ):
        """Initializes a new Decoder instance.

        Args:
            in_channels: Number of input channels.
            out_channels: Number of output channels.
            hidden_channels: Number of channels in (non residual block) hidden layers.
            n_residual_blocks: Number of residual blocks in each residual stack.
            residual_channels: Number of hidden channels in residual blocks.
            stride: Stride to use in the upsampling (i.e. transpose) convolutions. Must
                be even.
        """
        super().__init__()

        assert stride % 2 == 0, '"stride" must be even.'

        net = [
            nn.Conv2d(
                in_channels=in_channels,
                out_channels=hidden_channels,
                kernel_size=3,
                padding=1,
            ),
            ResidualStack(
                n_channels=hidden_channels,
                hidden_channels=residual_channels,
                n_residual_blocks=n_residual_blocks,
            ),
        ]
        for i in range(stride // 2):
            first, last = 0, stride // 2 - 1
            in_c = hidden_channels if i == first else hidden_channels // 2
            out_c = hidden_channels // 2 if i < last else out_channels
            net.append(
                nn.ConvTranspose2d(
                    in_channels=in_c,
                    out_channels=out_c,
                    kernel_size=4,
                    stride=2,
                    padding=1,
                )
            )
            if i < last:
                net.append(nn.ReLU())
        self._net = nn.Sequential(*net)

    def forward(self, x):
        return self._net(x)


class Quantizer(nn.Module):
    """Wraps a VectorQuantizer to handle input with arbitrary channels."""

    def __init__(self, in_channels, n_embeddings, embedding_dim):
        """Initializes a new Quantizer instance.

        Args:
            in_channels: Number of input channels.
            n_embeddings: Number of VectorQuantizer embeddings.
            embedding_dim: VectorQuantizer embedding dimension.
        """
        super().__init__()
        self._net = nn.Sequential(
            nn.Conv2d(
                in_channels=in_channels, out_channels=embedding_dim, kernel_size=1
            ),
            pg_nn.VectorQuantizer(n_embeddings, embedding_dim),
        )

    def forward(self, x):
        return self._net(x)

<filename>pytorch_generative/models/vae/vaes.py
"""Common modules used by Variational Autoencoders."""

import numpy as np
import torch
from torch import nn

from pytorch_generative import nn as pg_nn


@torch.jit.script
def to_var(log_std):
    """Returns the variance given the log standard deviation."""
    return log_std.exp().pow(2)


@torch.jit.script
def unit_gaussian_kl_div(mean, log_std):
    """Returns `KL(p || N(0, 1))` where `p` is a Gaussian with diagonal covariance."""
    return -0.5 * (1 + 2 * log_std - to_var(log_std) - mean ** 2)


@torch.jit.script
def gaussian_kl_div(p_mean, p_log_std, q_mean, q_log_std):
    """Returns `KL(p || q)` where `p` and `q` are Gaussians with diagonal covariance."""
    mean_delta, log_std_delta = (p_mean - q_mean) ** 2, q_log_std - p_log_std
    p_var, q_var = to_var(p_log_std), 2 * to_var(q_log_std)
    return -0.5 + log_std_delta + (p_var + mean_delta) / q_var


@torch.jit.script
def sample_from_gaussian(mu, log_sig):
    """Returns a sample from a Gaussian with diagonal covariance."""
    return mu + log_sig.exp() * torch.randn_like(log_sig)


@torch.jit.script
def _unflatten_tril(x):
    """Unflattens a vector into a lower triangular matrix of shape `dim x dim`."""
    n, dim = x.shape
    idxs = torch.tril_indices(dim, dim)
    tril = torch.zeros(n, dim, dim)
    tril[:, idxs[0, :], idxs[1, :]] = x
    return tril


@torch.jit.script
def gaussian_log_prob(x, mu, chol_sig):
    """Returns the log likelihood of `x` under a Gaussian with full covariance.

    The covariance is assumed to be positive-definite and to have been decomposed into
    the Cholesky factor `chol_sig`, i.e. `sig = chol_sig @ chol_sig.T`.

    Args:
        x: The input for which to compute the log likelihood.
        mu: The mean of the multivariate Gaussian.
        chol_sig: The flattened lower-triangular Cholesky factor of the covariance.
    Returns:
        The log likelihood.
    """
    dim = x.shape[0]
    chol_sig = _unflatten_tril(chol_sig)
    sig = chol_sig @ chol_sig.T
    const = -0.5 * dim * torch.log(torch.tensor(2 * np.pi))
    log_det = -0.5 * torch.logdet(x)
    exp = -0.5 * ((x - mu).T @ sig.inverse() @ (x - mu))
    return const + log_det + exp


class ResidualBlock(nn.Module):
    """A simple residual block."""

    def __init__(self, n_channels, hidden_channels):
        """Initializes a new ResidualBlock instance.

        Args:
            n_channels: Number of input and output channels.
            hidden_channels: Number of hidden channels.
        """
        super().__init__()
        self._net = nn.Sequential(
            nn.ReLU(),
            nn.Conv2d(
                in_channels=n_channels,
                out_channels=hidden_channels,
                kernel_size=3,
                padding=1,
            ),
            nn.ReLU(),
            nn.Conv2d(
                in_channels=hidden_channels, out_channels=n_channels, kernel_size=1
            ),
        )

    def forward(self, x):
        return x + self._net(x)


class ResidualStack(nn.Module):
    """A stack of multiple ResidualBlocks."""

    def __init__(self, n_channels, hidden_channels, n_residual_blocks=1):
        """Initializes a new ResidualStack instance.

        Args:
            n_channels: Number of input and output channels.
            hidden_channels: Number of hidden channels.
            n_residual_blocks: Number of residual blocks in the stack.
        """
        super().__init__()
        self._net = nn.Sequential(
            *[
                ResidualBlock(n_channels, hidden_channels)
                for _ in range(n_residual_blocks)
            ]
            + [nn.ReLU()]
        )

    def forward(self, x):
        return self._net(x)


class Encoder(nn.Module):
    """A feedforward encoder which downsamples its input."""

    def __init__(
        self,
        in_channels,
        out_channels,
        hidden_channels,
        n_residual_blocks,
        residual_channels,
        stride,
    ):
        """Initializes a new Encoder instance.

        Args:
            in_channels: Number of input channels.
            out_channels: Number of output channels.
            hidden_channels: Number of channels in non residual block hidden layers.
            n_residual_blocks: Number of residual blocks in each residual stack.
            residual_channels: Number of hidden channels in residual blocks.
            stride: Stride to use in the downsampling convolutions. Must be even.
        """
        super().__init__()
        assert stride % 2 == 0, '"stride" must be even.'

        net = []
        for i in range(stride // 2):
            first, last = 0, stride // 2 - 1
            in_c = in_channels if i == first else hidden_channels // 2
            out_c = hidden_channels // 2 if i < last else hidden_channels
            net.append(
                nn.Conv2d(
                    in_channels=in_c,
                    out_channels=out_c,
                    kernel_size=4,
                    stride=2,
                    padding=1,
                )
            )
            net.append(nn.ReLU())
        net.append(
            ResidualStack(
                n_channels=hidden_channels,
                hidden_channels=residual_channels,
                n_residual_blocks=n_residual_blocks,
            )
        )
        net.append(
            nn.Conv2d(
                in_channels=hidden_channels,
                out_channels=out_channels,
                kernel_size=3,
                padding=1,
            )
        )
        self._net = nn.Sequential(*net)

    def forward(self, x):
        return self._net(x)


class Decoder(nn.Module):
    """A feedforward encoder which upsamples its input."""

    def __init__(
        self,
        in_channels,
        out_channels,
        hidden_channels,
        n_residual_blocks,
        residual_channels,
        stride,
    ):
        """Initializes a new Decoder instance.

        Args:
            in_channels: Number of input channels.
            out_channels: Number of output channels.
            hidden_channels: Number of channels in (non residual block) hidden layers.
            n_residual_blocks: Number of residual blocks in each residual stack.
            residual_channels: Number of hidden channels in residual blocks.
            stride: Stride to use in the upsampling (i.e. transpose) convolutions. Must
                be even.
        """
        super().__init__()

        assert stride % 2 == 0, '"stride" must be even.'

        net = [
            nn.Conv2d(
                in_channels=in_channels,
                out_channels=hidden_channels,
                kernel_size=3,
                padding=1,
            ),
            ResidualStack(
                n_channels=hidden_channels,
                hidden_channels=residual_channels,
                n_residual_blocks=n_residual_blocks,
            ),
        ]
        for i in range(stride // 2):
            first, last = 0, stride // 2 - 1
            in_c = hidden_channels if i == first else hidden_channels // 2
            out_c = hidden_channels // 2 if i < last else out_channels
            net.append(
                nn.ConvTranspose2d(
                    in_channels=in_c,
                    out_channels=out_c,
                    kernel_size=4,
                    stride=2,
                    padding=1,
                )
            )
            if i < last:
                net.append(nn.ReLU())
        self._net = nn.Sequential(*net)

    def forward(self, x):
        return self._net(x)


class Quantizer(nn.Module):
    """Wraps a VectorQuantizer to handle input with arbitrary channels."""

    def __init__(self, in_channels, n_embeddings, embedding_dim):
        """Initializes a new Quantizer instance.

        Args:
            in_channels: Number of input channels.
            n_embeddings: Number of VectorQuantizer embeddings.
            embedding_dim: VectorQuantizer embedding dimension.
        """
        super().__init__()
        self._net = nn.Sequential(
            nn.Conv2d(
                in_channels=in_channels, out_channels=embedding_dim, kernel_size=1
            ),
            pg_nn.VectorQuantizer(n_embeddings, embedding_dim),
        )

    def forward(self, x):
        return self._net(x)

Common modules used by Variational Autoencoders. Returns the variance given the log standard deviation. Returns `KL(p || N(0, 1))` where `p` is a Gaussian with diagonal covariance. Returns `KL(p || q)` where `p` and `q` are Gaussians with diagonal covariance. Returns a sample from a Gaussian with diagonal covariance. Unflattens a vector into a lower triangular matrix of shape `dim x dim`. Returns the log likelihood of `x` under a Gaussian with full covariance.      The covariance is assumed to be positive-definite and to have been decomposed into     the Cholesky factor `chol_sig`, i.e. `sig = chol_sig @ chol_sig.T`.      Args:         x: The input for which to compute the log likelihood.         mu: The mean of the multivariate Gaussian.         chol_sig: The flattened lower-triangular Cholesky factor of the covariance.     Returns:         The log likelihood. A simple residual block. Initializes a new ResidualBlock instance.          Args:             n_channels: Number of input and output channels.             hidden_channels: Number of hidden channels. A stack of multiple ResidualBlocks. Initializes a new ResidualStack instance.          Args:             n_channels: Number of input and output channels.             hidden_channels: Number of hidden channels.             n_residual_blocks: Number of residual blocks in the stack. A feedforward encoder which downsamples its input. Initializes a new Encoder instance.          Args:             in_channels: Number of input channels.             out_channels: Number of output channels.             hidden_channels: Number of channels in non residual block hidden layers.             n_residual_blocks: Number of residual blocks in each residual stack.             residual_channels: Number of hidden channels in residual blocks.             stride: Stride to use in the downsampling convolutions. Must be even. A feedforward encoder which upsamples its input. Initializes a new Decoder instance.          Args:             in_channels: Number of input channels.             out_channels: Number of output channels.             hidden_channels: Number of channels in (non residual block) hidden layers.             n_residual_blocks: Number of residual blocks in each residual stack.             residual_channels: Number of hidden channels in residual blocks.             stride: Stride to use in the upsampling (i.e. transpose) convolutions. Must                 be even. Wraps a VectorQuantizer to handle input with arbitrary channels. Initializes a new Quantizer instance.          Args:             in_channels: Number of input channels.             n_embeddings: Number of VectorQuantizer embeddings.             embedding_dim: VectorQuantizer embedding dimension.