mirror of
https://github.com/zebrajr/pytorch.git
synced 2025-12-07 00:21:07 +01:00
This is a new version of #15648 based on the latest master branch. Unlike the previous PR where I fixed a lot of the doctests in addition to integrating xdoctest, I'm going to reduce the scope here. I'm simply going to integrate xdoctest, and then I'm going to mark all of the failing tests as "SKIP". This will let xdoctest run on the dashboards, provide some value, and still let the dashboards pass. I'll leave fixing the doctests themselves to another PR. In my initial commit, I do the bare minimum to get something running with failing dashboards. The few tests that I marked as skip are causing segfaults. Running xdoctest results in 293 failed, 201 passed tests. The next commits will be to disable those tests. (unfortunately I don't have a tool that will insert the `#xdoctest: +SKIP` directive over every failing test, so I'm going to do this mostly manually.) Fixes https://github.com/pytorch/pytorch/issues/71105 @ezyang Pull Request resolved: https://github.com/pytorch/pytorch/pull/82797 Approved by: https://github.com/ezyang
98 lines
3.6 KiB
Python
98 lines
3.6 KiB
Python
import math
|
|
|
|
import torch
|
|
from torch._six 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-deterinistic")
|
|
>>> 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
|
|
|
|
@property
|
|
def mean(self):
|
|
m = self.loc.clone(memory_format=torch.contiguous_format)
|
|
m[self.df <= 1] = nan
|
|
return m
|
|
|
|
@property
|
|
def mode(self):
|
|
return self.loc
|
|
|
|
@property
|
|
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., scale=1., 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(StudentT, self).__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.)))
|
|
return -0.5 * (self.df + 1.) * torch.log1p(y**2. / 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)
|