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