pytorch
35 строк · 1002.0 Байт
1# mypy: allow-untyped-defs
2from torch.distributions import constraints3from torch.distributions.gamma import Gamma4
5
6__all__ = ["Chi2"]7
8
9class Chi2(Gamma):10r"""11Creates a Chi-squared distribution parameterized by shape parameter :attr:`df`.
12This is exactly equivalent to ``Gamma(alpha=0.5*df, beta=0.5)``
13
14Example::
15
16>>> # xdoctest: +IGNORE_WANT("non-deterministic")
17>>> m = Chi2(torch.tensor([1.0]))
18>>> m.sample() # Chi2 distributed with shape df=1
19tensor([ 0.1046])
20
21Args:
22df (float or Tensor): shape parameter of the distribution
23"""
24arg_constraints = {"df": constraints.positive}25
26def __init__(self, df, validate_args=None):27super().__init__(0.5 * df, 0.5, validate_args=validate_args)28
29def expand(self, batch_shape, _instance=None):30new = self._get_checked_instance(Chi2, _instance)31return super().expand(batch_shape, new)32
33@property34def df(self):35return self.concentration * 236