From 6f4409073f9fad32016f18b3b083f959d7600a2b Mon Sep 17 00:00:00 2001 From: min-jean-cho Date: Wed, 15 Nov 2023 03:48:59 +0000 Subject: [PATCH] [doc] two diff meanings of rv generated by torch.tensor.geometric_ and torch.distributions.geometric.Geometric (#113183) The meaning of random variables generated by `torch.tensor.geometric_` and `torch.distributions.geometric.Geometric` are different, and they are defined by two different PMFs. Inform the user, so the user can choose their desired one. Background: https://github.com/pytorch/pytorch/pull/37984#issuecomment-630336511 Pull Request resolved: https://github.com/pytorch/pytorch/pull/113183 Approved by: https://github.com/albanD --- torch/_tensor_docs.py | 6 +++++- torch/distributions/geometric.py | 11 ++++++++--- 2 files changed, 13 insertions(+), 4 deletions(-) diff --git a/torch/_tensor_docs.py b/torch/_tensor_docs.py index ad815f585f17b..eb15749d40463 100644 --- a/torch/_tensor_docs.py +++ b/torch/_tensor_docs.py @@ -2114,8 +2114,12 @@ def add_docstr_all(method, docstr): .. math:: - P(X=k) = (1 - p)^{k - 1} p + P(X=k) = (1 - p)^{k - 1} p, k = 1, 2, ... +.. note:: + :func:`torch.Tensor.geometric_` `k`-th trial is the first success hence draws samples in :math:`\{1, 2, \ldots\}`, whereas + :func:`torch.distributions.geometric.Geometric` :math:`(k+1)`-th trial is the first success + hence draws samples in :math:`\{0, 1, \ldots\}`. """, ) diff --git a/torch/distributions/geometric.py b/torch/distributions/geometric.py index aba5373dc856d..0bf2f3dbacc67 100644 --- a/torch/distributions/geometric.py +++ b/torch/distributions/geometric.py @@ -18,10 +18,15 @@ class Geometric(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`). + .. math:: + + P(X=k) = (1-p)^{k} p, k = 0, 1, ... + + .. note:: + :func:`torch.distributions.geometric.Geometric` :math:`(k+1)`-th trial is the first success + hence draws samples in :math:`\{0, 1, \ldots\}`, whereas + :func:`torch.Tensor.geometric_` `k`-th trial is the first success hence draws samples in :math:`\{1, 2, \ldots\}`. Example::