【Hackathon 5th No.9】add multigammaln api -part (#57599)

GreatV · sunzhongkai588 · web-flow · commit 4eb03dca8fde · 2023-11-17T17:16:30.000+08:00
* add multigammaln api

* fix PR-CI-Static-Check

* fix PR-CI-Static-Check

* fix PR-CI-Static-Check

* update

* fix bug

* add test_inplace

* fix bug

* fix bug

* fixed

* bug fixed

* update

* fix bug

* Update python/paddle/tensor/math.py

Co-authored-by: zachary sun &lt;70642955+sunzhongkai588@users.noreply.github.com&gt;

---------

Co-authored-by: zachary sun &lt;70642955+sunzhongkai588@users.noreply.github.com&gt;
diff --git a/python/paddle/__init__.py b/python/paddle/__init__.py
@@ -318,6 +318,8 @@
     square_,
     stanh,
     sum,
+    multigammaln,
+    multigammaln_,
     nan_to_num,
     nan_to_num_,
     nansum,
@@ -887,6 +889,8 @@
     'renorm_',
     'take_along_axis',
     'put_along_axis',
+    'multigammaln',
+    'multigammaln_',
     'nan_to_num',
     'nan_to_num_',
     'heaviside',
diff --git a/python/paddle/tensor/__init__.py b/python/paddle/tensor/__init__.py
@@ -227,6 +227,8 @@
 from .math import square  # noqa: F401
 from .math import stanh  # noqa: F401
 from .math import sum  # noqa: F401
+from .math import multigammaln  # noqa: F401
+from .math import multigammaln_  # noqa: F401
 from .math import nan_to_num  # noqa: F401
 from .math import nan_to_num_  # noqa: F401
 from .math import nansum  # noqa: F401
@@ -473,6 +475,8 @@
     'square',
     'stanh',
     'sum',
+    'multigammaln',
+    'multigammaln_',
     'nan_to_num',
     'nan_to_num_',
     'hypot',
diff --git a/python/paddle/tensor/math.py b/python/paddle/tensor/math.py
@@ -15,6 +15,7 @@
 math functions
 """
 
+import math
 
 import numpy as np
 
@@ -5070,6 +5071,57 @@ def lgamma_(x, name=None):
         return _C_ops.lgamma_(x)
 
 
+def multigammaln(x, p, name=None):
+    """
+    This function computes the log of multivariate gamma, also sometimes called the generalized gamma.
+
+    Args:
+        x (Tensor): Input Tensor. Must be one of the following types: float16, float32, float64, uint16.
+        p (int): The dimension of the space of integration.
+        name (str, optional): Name for the operation (optional, default is None). For more information, please refer to :ref:`api_guide_Name`.
+
+    Returns:
+        out (Tensor): The values of the log multivariate gamma at the given tensor x.
+
+    Examples:
+        .. code-block:: python
+
+            >>> import paddle
+
+            >>> x = paddle.to_tensor([2.5, 3.5, 4, 6.5, 7.8, 10.23, 34.25])
+            >>> p = 2
+            >>> out = paddle.multigammaln(x, p)
+            >>> print(out)
+            Tensor(shape=[7], dtype=float32, place=Place(cpu), stop_gradient=True,
+                [0.85704780  , 2.46648574  , 3.56509781  , 11.02241898 , 15.84497833 ,
+                    26.09257698 , 170.68318176])
+    """
+    assert p >= 1, (
+        "The p must be greater than or equal to 1, "
+        "But received p is %s.\n" % p
+    )
+    c = 0.25 * p * (p - 1) * math.log(math.pi)
+    b = 0.5 * paddle.arange(start=(1 - p), end=1, step=1, dtype=x.dtype)
+    return paddle.sum(paddle.lgamma(x.unsqueeze(-1) + b), axis=-1) + c
+
+
+@inplace_apis_in_dygraph_only
+def multigammaln_(x, p, name=None):
+    r"""
+    Inplace version of ``multigammaln_`` API, the output Tensor will be inplaced with input ``x``.
+    Please refer to :ref:`api_paddle_multigammaln`.
+    """
+    assert p >= 1, (
+        "The p must be greater than or equal to 1, "
+        "But received p is %s.\n" % p
+    )
+    c = 0.25 * p * (p - 1) * math.log(math.pi)
+    c = paddle.to_tensor(c, dtype=x.dtype)
+    b = 0.5 * paddle.arange(start=(1 - p), end=1, step=1, dtype=x.dtype)
+    paddle.assign((x.unsqueeze(-1) + b).lgamma_().sum(-1).add_(c), x)
+    return x
+
+
 def neg(x, name=None):
     """
     This function computes the negative of the Tensor elementwisely.
diff --git a/test/legacy_test/test_inplace.py b/test/legacy_test/test_inplace.py
@@ -837,6 +837,22 @@ def non_inplace_api_processing(self, var):
         return paddle.digamma(var)
 
 
+class TestDygraphInplaceMutilgammaln(TestDygraphInplaceWithContinuous):
+    def init_data(self):
+        self.input_var_numpy = np.random.rand(10, 20).astype('float32') + 1.0
+        self.dtype = "float32"
+        self.p = 2
+
+    def inplace_api_processing(self, var):
+        return paddle.multigammaln_(var, self.p)
+
+    def non_inplace_api_processing(self, var):
+        return paddle.multigammaln(var, self.p)
+
+    def test_leaf_inplace_var_error(self):
+        pass
+
+
 class TestDygraphInplaceNeg(TestDygraphInplaceWithContinuous):
     def inplace_api_processing(self, var):
         return paddle.neg_(var)
diff --git a/test/legacy_test/test_multigammaln.py b/test/legacy_test/test_multigammaln.py
@@ -0,0 +1,111 @@
+#  Copyright (c) 2023 PaddlePaddle Authors. All Rights Reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#    https://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+import unittest
+
+import numpy as np
+from scipy import special
+
+import paddle
+
+
+def ref_multigammaln(x, p):
+    return special.multigammaln(x, p)
+
+
+def ref_multigammaln_grad(x, p):
+    def single_multigammaln_grad(x, p):
+        return special.psi(x - 0.5 * np.arange(0, p)).sum()
+
+    vectorized_multigammaln_grad = np.vectorize(single_multigammaln_grad)
+    return vectorized_multigammaln_grad(x, p)
+
+
+class TestMultigammalnAPI(unittest.TestCase):
+    def setUp(self):
+        np.random.seed(1024)
+        self.x = np.random.rand(10, 20).astype('float32') + 1.0
+        self.p = 2
+        self.init_input()
+        self.place = (
+            paddle.CUDAPlace(0)
+            if paddle.is_compiled_with_cuda()
+            else paddle.CPUPlace()
+        )
+
+    def init_input(self):
+        pass
+
+    def test_static_api(self):
+        paddle.enable_static()
+        with paddle.static.program_guard(paddle.static.Program()):
+            x = paddle.static.data('x', self.x.shape, dtype=self.x.dtype)
+            out = paddle.multigammaln(x, self.p)
+            exe = paddle.static.Executor(self.place)
+            res = exe.run(
+                feed={
+                    'x': self.x,
+                },
+                fetch_list=[out],
+            )
+            out_ref = ref_multigammaln(self.x, self.p)
+            np.testing.assert_allclose(out_ref, res[0], rtol=1e-6, atol=1e-6)
+
+    def test_dygraph_api(self):
+        paddle.disable_static(self.place)
+        x = paddle.to_tensor(self.x)
+        out = paddle.multigammaln(x, self.p)
+        out_ref = ref_multigammaln(self.x, self.p)
+        np.testing.assert_allclose(out_ref, out.numpy(), rtol=1e-6, atol=1e-6)
+        paddle.enable_static()
+
+
+class TestMultigammalnAPICase1(TestMultigammalnAPI):
+    def init_input(self):
+        self.x = np.random.rand(10, 20).astype('float64') + 1.0
+
+
+class TestMultigammalnGrad(unittest.TestCase):
+    def setUp(self):
+        np.random.seed(1024)
+        self.dtype = 'float32'
+        self.x = np.array([2, 3, 4, 5, 6, 7, 8]).astype(dtype=self.dtype)
+        self.p = 3
+        self.place = (
+            paddle.CUDAPlace(0)
+            if paddle.is_compiled_with_cuda()
+            else paddle.CPUPlace()
+        )
+
+    def test_backward(self):
+        expected_x_grad = ref_multigammaln_grad(self.x, self.p)
+        paddle.disable_static(self.place)
+        x = paddle.to_tensor(self.x, dtype=self.dtype, place=self.place)
+        x.stop_gradient = False
+        out = x.multigammaln(self.p)
+        loss = out.sum()
+        loss.backward()
+
+        np.testing.assert_allclose(
+            x.grad.numpy().astype('float32'),
+            expected_x_grad,
+            rtol=1e-6,
+            atol=1e-6,
+        )
+        paddle.enable_static()
+
+
+if __name__ == '__main__':
+    paddle.enable_static()
+    unittest.main()
