def __init__(self, rate, factor):
assert (factor <= 1).all()
self.rate, self.factor = broadcast_all(rate, factor)
propose = Exponential(self.factor * self.rate)
log_scale = self.factor.log()
super(RejectionExponential,
self).__init__(propose, self.log_prob_accept, log_scale)
def sample(self, sample_shape=()):
denom = torch.arange(0.5,
self.num_gamma_variates,
device=self.prototype.device).pow(2.0)
ones = self.prototype.new_ones((self.num_gamma_variates))
x = Exponential(ones).sample(self.batch_shape + sample_shape)
x = (x / denom).sum(-1)
return torch.clamp(x * (0.5 / math.pi**2), max=self.truncation_point)