[docs]classGeometric(Distribution):r""" Creates a Geometric distribution parameterized by :attr:`probs`, where :attr:`probs` is the probability of success of Bernoulli trials. It represents the probability that in :math:`k + 1` Bernoulli trials, the first :math:`k` trials failed, before seeing a success. Samples are non-negative integers [0, :math:`\inf`). Example:: >>> # xdoctest: +IGNORE_WANT("non-deterinistic") >>> m = Geometric(torch.tensor([0.3])) >>> m.sample() # underlying Bernoulli has 30% chance 1; 70% chance 0 tensor([ 2.]) Args: probs (Number, Tensor): the probability of sampling `1`. Must be in range (0, 1] logits (Number, Tensor): the log-odds of sampling `1`. """arg_constraints={"probs":constraints.unit_interval,"logits":constraints.real}support=constraints.nonnegative_integerdef__init__(self,probs=None,logits=None,validate_args=None):if(probsisNone)==(logitsisNone):raiseValueError("Either `probs` or `logits` must be specified, but not both.")ifprobsisnotNone:(self.probs,)=broadcast_all(probs)else:(self.logits,)=broadcast_all(logits)probs_or_logits=probsifprobsisnotNoneelselogitsifisinstance(probs_or_logits,Number):batch_shape=torch.Size()else:batch_shape=probs_or_logits.size()super().__init__(batch_shape,validate_args=validate_args)ifself._validate_argsandprobsisnotNone:# Add an extra check beyond unit_intervalvalue=self.probsvalid=value>0ifnotvalid.all():invalid_value=value.data[~valid]raiseValueError("Expected parameter probs "f"({type(value).__name__} of shape {tuple(value.shape)}) "f"of distribution {repr(self)} "f"to be positive but found invalid values:\n{invalid_value}")
[docs]defsample(self,sample_shape=torch.Size()):shape=self._extended_shape(sample_shape)tiny=torch.finfo(self.probs.dtype).tinywithtorch.no_grad():iftorch._C._get_tracing_state():# [JIT WORKAROUND] lack of support for .uniform_()u=torch.rand(shape,dtype=self.probs.dtype,device=self.probs.device)u=u.clamp(min=tiny)else:u=self.probs.new(shape).uniform_(tiny,1)return(u.log()/(-self.probs).log1p()).floor()
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.