Source code for torch.distributions.gumbel

import math
from numbers import Number

import torch
from torch.distributions import constraints
from torch.distributions.transformed_distribution import TransformedDistribution
from torch.distributions.transforms import AffineTransform, ExpTransform
from torch.distributions.uniform import Uniform
from torch.distributions.utils import broadcast_all, euler_constant

__all__ = ["Gumbel"]

[docs]class Gumbel(TransformedDistribution): r""" Samples from a Gumbel Distribution. Examples:: >>> # xdoctest: +IGNORE_WANT("non-deterministic") >>> m = Gumbel(torch.tensor([1.0]), torch.tensor([2.0])) >>> m.sample() # sample from Gumbel distribution with loc=1, scale=2 tensor([ 1.0124]) Args: loc (float or Tensor): Location parameter of the distribution scale (float or Tensor): Scale parameter of the distribution """ arg_constraints = {"loc": constraints.real, "scale": constraints.positive} support = constraints.real def __init__(self, loc, scale, validate_args=None): self.loc, self.scale = broadcast_all(loc, scale) finfo = torch.finfo(self.loc.dtype) if isinstance(loc, Number) and isinstance(scale, Number): base_dist = Uniform(finfo.tiny, 1 - finfo.eps, validate_args=validate_args) else: base_dist = Uniform( torch.full_like(self.loc, finfo.tiny), torch.full_like(self.loc, 1 - finfo.eps), validate_args=validate_args, ) transforms = [ ExpTransform().inv, AffineTransform(loc=0, scale=-torch.ones_like(self.scale)), ExpTransform().inv, AffineTransform(loc=loc, scale=-self.scale), ] super().__init__(base_dist, transforms, validate_args=validate_args)
[docs] def expand(self, batch_shape, _instance=None): new = self._get_checked_instance(Gumbel, _instance) new.loc = self.loc.expand(batch_shape) new.scale = self.scale.expand(batch_shape) return super().expand(batch_shape, _instance=new)
# Explicitly defining the log probability function for Gumbel due to precision issues
[docs] def log_prob(self, value): if self._validate_args: self._validate_sample(value) y = (self.loc - value) / self.scale return (y - y.exp()) - self.scale.log()
@property def mean(self): return self.loc + self.scale * euler_constant @property def mode(self): return self.loc @property def stddev(self): return (math.pi / math.sqrt(6)) * self.scale @property def variance(self): return self.stddev.pow(2)
[docs] def entropy(self): return self.scale.log() + (1 + euler_constant)


Access comprehensive developer documentation for PyTorch

View Docs


Get in-depth tutorials for beginners and advanced developers

View Tutorials


Find development resources and get your questions answered

View Resources