Spaces:
Sleeping
Sleeping
| import math | |
| import torch | |
| from torch import inf, nan | |
| from torch.distributions import Chi2, constraints | |
| from torch.distributions.distribution import Distribution | |
| from torch.distributions.utils import _standard_normal, broadcast_all | |
| __all__ = ["StudentT"] | |
| class StudentT(Distribution): | |
| r""" | |
| Creates a Student's t-distribution parameterized by degree of | |
| freedom :attr:`df`, mean :attr:`loc` and scale :attr:`scale`. | |
| Example:: | |
| >>> # xdoctest: +IGNORE_WANT("non-deterministic") | |
| >>> m = StudentT(torch.tensor([2.0])) | |
| >>> m.sample() # Student's t-distributed with degrees of freedom=2 | |
| tensor([ 0.1046]) | |
| Args: | |
| df (float or Tensor): degrees of freedom | |
| loc (float or Tensor): mean of the distribution | |
| scale (float or Tensor): scale of the distribution | |
| """ | |
| arg_constraints = { | |
| "df": constraints.positive, | |
| "loc": constraints.real, | |
| "scale": constraints.positive, | |
| } | |
| support = constraints.real | |
| has_rsample = True | |
| def mean(self): | |
| m = self.loc.clone(memory_format=torch.contiguous_format) | |
| m[self.df <= 1] = nan | |
| return m | |
| def mode(self): | |
| return self.loc | |
| def variance(self): | |
| m = self.df.clone(memory_format=torch.contiguous_format) | |
| m[self.df > 2] = ( | |
| self.scale[self.df > 2].pow(2) | |
| * self.df[self.df > 2] | |
| / (self.df[self.df > 2] - 2) | |
| ) | |
| m[(self.df <= 2) & (self.df > 1)] = inf | |
| m[self.df <= 1] = nan | |
| return m | |
| def __init__(self, df, loc=0.0, scale=1.0, validate_args=None): | |
| self.df, self.loc, self.scale = broadcast_all(df, loc, scale) | |
| self._chi2 = Chi2(self.df) | |
| batch_shape = self.df.size() | |
| super().__init__(batch_shape, validate_args=validate_args) | |
| def expand(self, batch_shape, _instance=None): | |
| new = self._get_checked_instance(StudentT, _instance) | |
| batch_shape = torch.Size(batch_shape) | |
| new.df = self.df.expand(batch_shape) | |
| new.loc = self.loc.expand(batch_shape) | |
| new.scale = self.scale.expand(batch_shape) | |
| new._chi2 = self._chi2.expand(batch_shape) | |
| super(StudentT, new).__init__(batch_shape, validate_args=False) | |
| new._validate_args = self._validate_args | |
| return new | |
| def rsample(self, sample_shape=torch.Size()): | |
| # NOTE: This does not agree with scipy implementation as much as other distributions. | |
| # (see https://github.com/fritzo/notebooks/blob/master/debug-student-t.ipynb). Using DoubleTensor | |
| # parameters seems to help. | |
| # X ~ Normal(0, 1) | |
| # Z ~ Chi2(df) | |
| # Y = X / sqrt(Z / df) ~ StudentT(df) | |
| shape = self._extended_shape(sample_shape) | |
| X = _standard_normal(shape, dtype=self.df.dtype, device=self.df.device) | |
| Z = self._chi2.rsample(sample_shape) | |
| Y = X * torch.rsqrt(Z / self.df) | |
| return self.loc + self.scale * Y | |
| def log_prob(self, value): | |
| if self._validate_args: | |
| self._validate_sample(value) | |
| y = (value - self.loc) / self.scale | |
| Z = ( | |
| self.scale.log() | |
| + 0.5 * self.df.log() | |
| + 0.5 * math.log(math.pi) | |
| + torch.lgamma(0.5 * self.df) | |
| - torch.lgamma(0.5 * (self.df + 1.0)) | |
| ) | |
| return -0.5 * (self.df + 1.0) * torch.log1p(y**2.0 / self.df) - Z | |
| def entropy(self): | |
| lbeta = ( | |
| torch.lgamma(0.5 * self.df) | |
| + math.lgamma(0.5) | |
| - torch.lgamma(0.5 * (self.df + 1)) | |
| ) | |
| return ( | |
| self.scale.log() | |
| + 0.5 | |
| * (self.df + 1) | |
| * (torch.digamma(0.5 * (self.df + 1)) - torch.digamma(0.5 * self.df)) | |
| + 0.5 * self.df.log() | |
| + lbeta | |
| ) | |