pytorch

1
from numbers import Number
2

3
import torch
4
from torch import nan
5
from torch.distributions import constraints
6
from torch.distributions.distribution import Distribution
7
from torch.distributions.utils import broadcast_all
8

9
__all__ = ["Uniform"]
10

11

12
class Uniform(Distribution):
13
    r"""
14
    Generates uniformly distributed random samples from the half-open interval
15
    ``[low, high)``.
16

17
    Example::
18

19
        >>> m = Uniform(torch.tensor([0.0]), torch.tensor([5.0]))
20
        >>> m.sample()  # uniformly distributed in the range [0.0, 5.0)
21
        >>> # xdoctest: +SKIP
22
        tensor([ 2.3418])
23

24
    Args:
25
        low (float or Tensor): lower range (inclusive).
26
        high (float or Tensor): upper range (exclusive).
27
    """
28
    # TODO allow (loc,scale) parameterization to allow independent constraints.
29
    arg_constraints = {
30
        "low": constraints.dependent(is_discrete=False, event_dim=0),
31
        "high": constraints.dependent(is_discrete=False, event_dim=0),
32
    }
33
    has_rsample = True
34

35
    @property
36
    def mean(self):
37
        return (self.high + self.low) / 2
38

39
    @property
40
    def mode(self):
41
        return nan * self.high
42

43
    @property
44
    def stddev(self):
45
        return (self.high - self.low) / 12**0.5
46

47
    @property
48
    def variance(self):
49
        return (self.high - self.low).pow(2) / 12
50

51
    def __init__(self, low, high, validate_args=None):
52
        self.low, self.high = broadcast_all(low, high)
53

54
        if isinstance(low, Number) and isinstance(high, Number):
55
            batch_shape = torch.Size()
56
        else:
57
            batch_shape = self.low.size()
58
        super().__init__(batch_shape, validate_args=validate_args)
59

60
        if self._validate_args and not torch.lt(self.low, self.high).all():
61
            raise ValueError("Uniform is not defined when low>= high")
62

63
    def expand(self, batch_shape, _instance=None):
64
        new = self._get_checked_instance(Uniform, _instance)
65
        batch_shape = torch.Size(batch_shape)
66
        new.low = self.low.expand(batch_shape)
67
        new.high = self.high.expand(batch_shape)
68
        super(Uniform, new).__init__(batch_shape, validate_args=False)
69
        new._validate_args = self._validate_args
70
        return new
71

72
    @constraints.dependent_property(is_discrete=False, event_dim=0)
73
    def support(self):
74
        return constraints.interval(self.low, self.high)
75

76
    def rsample(self, sample_shape=torch.Size()):
77
        shape = self._extended_shape(sample_shape)
78
        rand = torch.rand(shape, dtype=self.low.dtype, device=self.low.device)
79
        return self.low + rand * (self.high - self.low)
80

81
    def log_prob(self, value):
82
        if self._validate_args:
83
            self._validate_sample(value)
84
        lb = self.low.le(value).type_as(self.low)
85
        ub = self.high.gt(value).type_as(self.low)
86
        return torch.log(lb.mul(ub)) - torch.log(self.high - self.low)
87

88
    def cdf(self, value):
89
        if self._validate_args:
90
            self._validate_sample(value)
91
        result = (value - self.low) / (self.high - self.low)
92
        return result.clamp(min=0, max=1)
93

94
    def icdf(self, value):
95
        result = value * (self.high - self.low) + self.low
96
        return result
97

98
    def entropy(self):
99
        return torch.log(self.high - self.low)
100