[docs]classNegativeBinomial(Distribution):r""" Creates a Negative Binomial distribution, i.e. distribution of the number of successful independent and identical Bernoulli trials before :attr:`total_count` failures are achieved. The probability of success of each Bernoulli trial is :attr:`probs`. Args: total_count (float or Tensor): non-negative number of negative Bernoulli trials to stop, although the distribution is still valid for real valued count probs (Tensor): Event probabilities of success in the half open interval [0, 1) logits (Tensor): Event log-odds for probabilities of success """arg_constraints={"total_count":constraints.greater_than_eq(0),"probs":constraints.half_open_interval(0.0,1.0),"logits":constraints.real,}support=constraints.nonnegative_integerdef__init__(self,total_count,probs=None,logits=None,validate_args=None):if(probsisNone)==(logitsisNone):raiseValueError("Either `probs` or `logits` must be specified, but not both.")ifprobsisnotNone:(self.total_count,self.probs,)=broadcast_all(total_count,probs)self.total_count=self.total_count.type_as(self.probs)else:(self.total_count,self.logits,)=broadcast_all(total_count,logits)self.total_count=self.total_count.type_as(self.logits)self._param=self.probsifprobsisnotNoneelseself.logitsbatch_shape=self._param.size()super().__init__(batch_shape,validate_args=validate_args)
def_new(self,*args,**kwargs):returnself._param.new(*args,**kwargs)@propertydefmean(self):returnself.total_count*torch.exp(self.logits)@propertydefmode(self):return((self.total_count-1)*self.logits.exp()).floor().clamp(min=0.0)@propertydefvariance(self):returnself.mean/torch.sigmoid(-self.logits)@lazy_propertydeflogits(self):returnprobs_to_logits(self.probs,is_binary=True)@lazy_propertydefprobs(self):returnlogits_to_probs(self.logits,is_binary=True)@propertydefparam_shape(self):returnself._param.size()@lazy_propertydef_gamma(self):# Note we avoid validating because self.total_count can be zero.returntorch.distributions.Gamma(concentration=self.total_count,rate=torch.exp(-self.logits),validate_args=False,)
[docs]deflog_prob(self,value):ifself._validate_args:self._validate_sample(value)log_unnormalized_prob=self.total_count*F.logsigmoid(-self.logits)+value*F.logsigmoid(self.logits)log_normalization=(-torch.lgamma(self.total_count+value)+torch.lgamma(1.0+value)+torch.lgamma(self.total_count))# The case self.total_count == 0 and value == 0 has probability 1 but# lgamma(0) is infinite. Handle this case separately using a function# that does not modify tensors in place to allow Jit compilation.log_normalization=log_normalization.masked_fill(self.total_count+value==0.0,0.0)returnlog_unnormalized_prob-log_normalization
Docs
Access comprehensive developer documentation for PyTorch
To analyze traffic and optimize your experience, we serve cookies on this site. By clicking or navigating, you agree to allow our usage of cookies. As the current maintainers of this site, Facebook’s Cookies Policy applies. Learn more, including about available controls: Cookies Policy.