Source code for torchdrug.layers.distribution

import math
from import Sequence

import torch
from torch import nn

[docs]class IndependentGaussian(nn.Module): """ Independent Gaussian distribution. Parameters: mu (Tensor): mean of shape :math:`(N,)` sigma2 (Tensor): variance of shape :math:`(N,)` learnable (bool, optional): learnable parameters or not """ def __init__(self, mu, sigma2, learnable=False): super(IndependentGaussian, self).__init__() if learnable: = nn.Parameter(torch.as_tensor(mu)) self.sigma2 = nn.Parameter(torch.as_tensor(sigma2)) else: self.register_buffer("mu", torch.as_tensor(mu)) self.register_buffer("sigma2", torch.as_tensor(sigma2)) self.dim = len(mu)
[docs] def forward(self, input): """ Compute the likelihood of input data. Parameters: input (Tensor): input data of shape :math:`(..., N)` """ log_likelihood = -0.5 * (math.log(2 * math.pi) + self.sigma2.log() + (input - ** 2 / self.sigma2) return log_likelihood
[docs] def sample(self, *size): """ Draw samples from the distribution. Parameters: size (tuple of int): shape of the samples """ if len(size) == 1 and isinstance(size[0], Sequence): size = size[0] size = list(size) + [self.dim] sample = torch.randn(size, * self.sigma2.sqrt() + return sample