From 44c20ce676c6ca2a8b04af72c80c36698ee34d60 Mon Sep 17 00:00:00 2001 From: Kushashwa Ravi Shrimali Date: Tue, 1 Jun 2021 23:02:36 -0700 Subject: [PATCH] Alias for `i0` to `special` namespace (#59141) Summary: See https://github.com/pytorch/pytorch/issues/50345 cc: mruberry kshitij12345 Pull Request resolved: https://github.com/pytorch/pytorch/pull/59141 Reviewed By: ngimel Differential Revision: D28784097 Pulled By: mruberry fbshipit-source-id: 9b61a21906ef337292686fd40e328502a79e6f09 --- aten/src/ATen/core/aten_interned_strings.h | 1 - aten/src/ATen/core/interned_strings.h | 2 ++ aten/src/ATen/native/UnaryOps.cpp | 4 ++++ aten/src/ATen/native/native_functions.yaml | 8 +++++++ docs/source/special.rst | 1 + torch/_torch_docs.py | 20 ++-------------- torch/csrc/api/include/torch/special.h | 24 +++++++++++++++---- torch/csrc/jit/passes/normalize_ops.cpp | 1 + torch/overrides.py | 1 + torch/special/__init__.py | 23 ++++++++++++++++++ .../_internal/common_methods_invocations.py | 1 + 11 files changed, 63 insertions(+), 23 deletions(-) diff --git a/aten/src/ATen/core/aten_interned_strings.h b/aten/src/ATen/core/aten_interned_strings.h index 819709bec66..0b6d6ea6312 100644 --- a/aten/src/ATen/core/aten_interned_strings.h +++ b/aten/src/ATen/core/aten_interned_strings.h @@ -376,7 +376,6 @@ _(aten, hspmm) \ _(aten, hsplit) \ _(aten, hstack) \ _(aten, hypot) \ -_(aten, i0) \ _(aten, i0_) \ _(aten, igamma) \ _(aten, igamma_) \ diff --git a/aten/src/ATen/core/interned_strings.h b/aten/src/ATen/core/interned_strings.h index 839e381df22..96a49f4426b 100644 --- a/aten/src/ATen/core/interned_strings.h +++ b/aten/src/ATen/core/interned_strings.h @@ -336,6 +336,8 @@ namespace c10 { _(aten, special_expm1) \ _(aten, exp2) \ _(aten, special_exp2) \ + _(aten, i0) \ + _(aten, special_i0) \ _(aten, special_i0e) \ _(aten, special_i1) \ _(aten, special_i1e) \ diff --git a/aten/src/ATen/native/UnaryOps.cpp b/aten/src/ATen/native/UnaryOps.cpp index 083ce7f73e8..07c0fc26383 100644 --- a/aten/src/ATen/native/UnaryOps.cpp +++ b/aten/src/ATen/native/UnaryOps.cpp @@ -412,6 +412,10 @@ Tensor special_erfc(const Tensor& self) { return self.erfc(); } Tensor& special_erfinv_out(const Tensor& self, Tensor& result) { return at::erfinv_out(result, self); } Tensor special_erfinv(const Tensor& self) { return self.erfinv(); } +// special_i0, alias for i0 +Tensor& special_i0_out(const Tensor& self, Tensor& result) { return at::i0_out(result, self); } +Tensor special_i0(const Tensor& self) { return self.i0(); } + namespace { inline Tensor calc_ndtr(const Tensor& self) { diff --git a/aten/src/ATen/native/native_functions.yaml b/aten/src/ATen/native/native_functions.yaml index f401c3ac809..7962368d2a9 100644 --- a/aten/src/ATen/native/native_functions.yaml +++ b/aten/src/ATen/native/native_functions.yaml @@ -9405,6 +9405,14 @@ dispatch: CompositeExplicitAutograd: special_xlog1py_out +- func: special_i0(Tensor self) -> Tensor + python_module: special + variants: function + +- func: special_i0.out(Tensor self, *, Tensor(a!) out) -> Tensor(a!) + python_module: special + variants: function + - func: special_i0e(Tensor self) -> Tensor python_module: special variants: function diff --git a/docs/source/special.rst b/docs/source/special.rst index a1e4b0afc06..39aa0640c95 100644 --- a/docs/source/special.rst +++ b/docs/source/special.rst @@ -26,6 +26,7 @@ Functions .. autofunction:: expm1 .. autofunction:: exp2 .. autofunction:: gammaln +.. autofunction:: i0 .. autofunction:: i0e .. autofunction:: i1 .. autofunction:: i1e diff --git a/torch/_torch_docs.py b/torch/_torch_docs.py index b3e243d70a1..9a49361a2ec 100644 --- a/torch/_torch_docs.py +++ b/torch/_torch_docs.py @@ -3866,24 +3866,8 @@ add_docstr(torch.i0, r""" i0(input, *, out=None) -> Tensor -Computes the zeroth order modified Bessel function of the first kind for each element of :attr:`input`. - -.. math:: - \text{out}_{i} = I_0(\text{input}_{i}) = \sum_{k=0}^{\infty} \frac{(\text{input}_{i}^2/4)^k}{(k!)^2} - -""" + r""" -Args: - input (Tensor): the input tensor - -Keyword args: - {out} - -Example:: - - >>> torch.i0(torch.arange(5, dtype=torch.float32)) - tensor([ 1.0000, 1.2661, 2.2796, 4.8808, 11.3019]) - -""".format(**common_args)) +Alias for :func:`torch.special.i0`. +""") add_docstr(torch.igamma, r""" diff --git a/torch/csrc/api/include/torch/special.h b/torch/csrc/api/include/torch/special.h index db7de24ea9f..d80a43981ae 100644 --- a/torch/csrc/api/include/torch/special.h +++ b/torch/csrc/api/include/torch/special.h @@ -182,6 +182,22 @@ inline Tensor& xlog1py_out(Tensor& result, const Tensor& self, const Scalar& oth return torch::special_xlog1py_out(result, self, other); } +/// Computes the zeroth order modified Bessel function of the first kind of input, elementwise +/// See https://pytorch.org/docs/master/special.html#torch.special.i0 +/// +/// Example: +/// ``` +/// auto t = torch::randn(128, dtype=kDouble); +/// torch::special::i0(t); +/// ``` +inline Tensor i0(const Tensor& self) { + return torch::special_i0(self); +} + +inline Tensor& i0_out(Tensor& result, const Tensor& self) { + return torch::special_i0_out(result, self); +} + /// Computes the area under the standard Gaussian probability density function, /// integrated from minus infinity to :attr:`input`, elementwise /// See https://pytorch.org/docs/master/special.html#torch.special.ndtr @@ -195,7 +211,7 @@ inline Tensor ndtr(const Tensor& self) { return torch::special_ndtr(self); } -inline Tensor ndtr_out(Tensor& result, const Tensor& self) { +inline Tensor& ndtr_out(Tensor& result, const Tensor& self) { return torch::special_ndtr_out(result, self); } @@ -211,7 +227,7 @@ inline Tensor i0e(const Tensor& self) { return torch::special_i0e(self); } -inline Tensor i0e_out(Tensor& result, const Tensor& self) { +inline Tensor& i0e_out(Tensor& result, const Tensor& self) { return torch::special_i0e_out(result, self); } @@ -227,7 +243,7 @@ inline Tensor i1(const Tensor& self) { return torch::special_i1(self); } -inline Tensor i1_out(Tensor& result, const Tensor& self) { +inline Tensor& i1_out(Tensor& result, const Tensor& self) { return torch::special_i1_out(result, self); } @@ -243,7 +259,7 @@ inline Tensor i1e(const Tensor& self) { return torch::special_i1e(self); } -inline Tensor i1e_out(Tensor& result, const Tensor& self) { +inline Tensor& i1e_out(Tensor& result, const Tensor& self) { return torch::special_i1e_out(result, self); } diff --git a/torch/csrc/jit/passes/normalize_ops.cpp b/torch/csrc/jit/passes/normalize_ops.cpp index 7e48a421ff5..eda56fe2277 100644 --- a/torch/csrc/jit/passes/normalize_ops.cpp +++ b/torch/csrc/jit/passes/normalize_ops.cpp @@ -117,6 +117,7 @@ const std::unordered_map& getOperatorAliasMap() { {aten::special_exp2, aten::exp2}, {aten::special_expm1, aten::expm1}, {aten::special_logit, aten::logit}, + {aten::special_i0, aten::i0}, {aten::orgqr, aten::linalg_householder_product}, {aten::special_gammaln, aten::lgamma}}; return alias_map; diff --git a/torch/overrides.py b/torch/overrides.py index bdb6d8c8cb0..f94c8756733 100644 --- a/torch/overrides.py +++ b/torch/overrides.py @@ -866,6 +866,7 @@ def get_testing_overrides() -> Dict[Callable, Callable]: torch.special.expm1: lambda input: -1, torch.special.expit: lambda input: -1, torch.special.gammaln: lambda input: -1, + torch.special.i0: lambda input: -1, torch.special.i0e: lambda input: -1, torch.special.i1: lambda input: -1, torch.special.i1e: lambda input: -1, diff --git a/torch/special/__init__.py b/torch/special/__init__.py index 7e690a418f8..d0aae87a0b2 100644 --- a/torch/special/__init__.py +++ b/torch/special/__init__.py @@ -272,6 +272,29 @@ Example:: tensor([2.7726, 2.1972, 1.3863]) """.format(**common_args)) +i0 = _add_docstr(_special.special_i0, + r""" +i0(input, *, out=None) -> Tensor + +Computes the zeroth order modified Bessel function of the first kind for each element of :attr:`input`. + +.. math:: + \text{out}_{i} = I_0(\text{input}_{i}) = \sum_{k=0}^{\infty} \frac{(\text{input}_{i}^2/4)^k}{(k!)^2} + +""" + r""" +Args: + input (Tensor): the input tensor + +Keyword args: + {out} + +Example:: + + >>> torch.i0(torch.arange(5, dtype=torch.float32)) + tensor([ 1.0000, 1.2661, 2.2796, 4.8808, 11.3019]) + +""".format(**common_args)) + i0e = _add_docstr(_special.special_i0e, r""" i0e(input, *, out=None) -> Tensor diff --git a/torch/testing/_internal/common_methods_invocations.py b/torch/testing/_internal/common_methods_invocations.py index 4fc5fdb25c8..c1c6483b6ee 100644 --- a/torch/testing/_internal/common_methods_invocations.py +++ b/torch/testing/_internal/common_methods_invocations.py @@ -4899,6 +4899,7 @@ op_db: List[OpInfo] = [ UnaryUfuncInfo('i0', ref=np_unary_ufunc_integer_promotion_wrapper( scipy.special.i0) if TEST_SCIPY else _NOTHING, + aliases=('special.i0',), decorators=(precisionOverride({torch.bfloat16: 3e-1, torch.float16: 5e-1}),), backward_dtypesIfCPU=floating_types(),