chore: import upstream snapshot with attribution

test/distribution/CMakeLists.txt

@@ -0,0 +1,14 @@
+file(
+  GLOB TEST_OPS
+  RELATIVE "${CMAKE_CURRENT_SOURCE_DIR}"
+  "test_*.py")
+string(REPLACE ".py" "" TEST_OPS "${TEST_OPS}")
+
+foreach(TEST_OP ${TEST_OPS})
+  py_test_modules(${TEST_OP} MODULES ${TEST_OP})
+endforeach()
+
+set_tests_properties(test_distribution_lkj_cholesky_static PROPERTIES TIMEOUT
+                                                                      120)
+
+set_pir_tests_properties()

test/distribution/config.py

@@ -0,0 +1,31 @@
+# Copyright (c) 2021 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
+#
+#     http://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 paddle
+
+DEVICES = [paddle.CPUPlace()]
+if paddle.is_compiled_with_cuda():
+    DEVICES.append(paddle.CUDAPlace(0))
+
+DEFAULT_DTYPE = 'float64'
+
+TEST_CASE_NAME = 'suffix'
+# All test case will use float64 for compare precision, refs:
+# https://github.com/PaddlePaddle/Paddle/wiki/Upgrade-OP-Precision-to-Float64
+RTOL = {
+    'float32': 1e-03,
+    'complex64': 1e-3,
+    'float64': 1e-5,
+    'complex128': 1e-5,
+}
+ATOL = {'float32': 0.0, 'complex64': 0, 'float64': 0.0, 'complex128': 0}

test/distribution/mock_data.py

@@ -0,0 +1,65 @@
+# Copyright (c) 2021 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
+#
+#     http://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 paddle
+
+
+class Exponential(paddle.distribution.ExponentialFamily):
+    """mock exponential distribution, which support computing entropy and
+    kl use bregman divergence
+    """
+
+    _mean_carrier_measure = 0
+
+    def __init__(self, rate):
+        self._rate = rate
+        super().__init__(batch_shape=rate.shape)
+
+    @property
+    def rate(self):
+        return self._rate
+
+    def entropy(self):
+        return 1.0 - paddle.log(self._rate)
+
+    @property
+    def _natural_parameters(self):
+        return (-self._rate,)
+
+    def _log_normalizer(self, x):
+        return -paddle.log(-x)
+
+
+class DummyExpFamily(paddle.distribution.ExponentialFamily):
+    """dummy class extend from exponential family"""
+
+    def __init__(self, *args):
+        pass
+
+    def entropy(self):
+        return 1.0
+
+    @property
+    def _natural_parameters(self):
+        return (1.0,)
+
+    def _log_normalizer(self, x):
+        return -paddle.log(-x)
+
+
+@paddle.distribution.register_kl(Exponential, Exponential)
+def _kl_exponential_exponential(p, q):
+    rate_ratio = q.rate / p.rate
+    t1 = -rate_ratio.log()
+    return t1 + rate_ratio - 1

test/distribution/parameterize.py

@@ -0,0 +1,240 @@
+# Copyright (c) 2021 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
+#
+#     http://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 collections
+import functools
+import inspect
+import re
+import sys
+from unittest import SkipTest
+
+import numpy as np
+from distribution import config
+
+TEST_CASE_NAME = 'suffix'
+
+
+def xrand(shape=(10, 10, 10), dtype=config.DEFAULT_DTYPE, min=1.0, max=10.0):
+    return (np.random.rand(*shape).astype(dtype)) * (max - min) + min
+
+
+def place(devices, key='place'):
+    def decorate(cls):
+        module = sys.modules[cls.__module__].__dict__
+        raw_classes = {
+            k: v for k, v in module.items() if k.startswith(cls.__name__)
+        }
+
+        for raw_name, raw_cls in raw_classes.items():
+            for d in devices:
+                test_cls = dict(raw_cls.__dict__)
+                test_cls.update({key: d})
+                new_name = raw_name + '.' + d.__class__.__name__
+                module[new_name] = type(new_name, (raw_cls,), test_cls)
+            del module[raw_name]
+        return cls
+
+    return decorate
+
+
+def parameterize_cls(fields, values=None, test_pir=False):
+    fields = [fields] if isinstance(fields, str) else fields
+    params = [dict(zip(fields, vals)) for vals in values]
+
+    def decorate(cls):
+        test_cls_module = sys.modules[cls.__module__].__dict__
+        for k, v in enumerate(params):
+            test_cls = dict(cls.__dict__)
+            test_cls.update(v)
+            test_cls["test_pir"] = False
+            name = cls.__name__ + str(k)
+            name = name + '.' + v.get('suffix') if v.get('suffix') else name
+            test_cls_module[name] = type(name, (cls,), test_cls)
+            if test_pir:
+                name = name + ".pir"
+                test_cls["test_pir"] = True
+                pir_type = type(name, (cls,), test_cls)
+                test_cls_module[name] = pir_type
+
+        for m in list(cls.__dict__):
+            if m.startswith("test"):
+                delattr(cls, m)
+        return cls
+
+    return decorate
+
+
+def parameterize_func(
+    input, name_func=None, doc_func=None, skip_on_empty=False
+):
+    name_func = name_func or default_name_func
+
+    def wrapper(f, instance=None):
+        frame_locals = inspect.currentframe().f_back.f_locals
+
+        parameters = input_as_callable(input)()
+
+        if not parameters:
+            if not skip_on_empty:
+                raise ValueError(
+                    "Parameters iterable is empty (hint: use "
+                    "`parameterized.expand([], skip_on_empty=True)` to skip "
+                    "this test when the input is empty)"
+                )
+            return functools.wraps(f)(skip_on_empty_helper)
+
+        digits = len(str(len(parameters) - 1))
+        for num, p in enumerate(parameters):
+            name = name_func(
+                f, "{num:0>{digits}}".format(digits=digits, num=num), p
+            )
+            # If the original function has patches applied by 'mock.patch',
+            # re-construct all patches on the just former decoration layer
+            # of param_as_standalone_func so as not to share
+            # patch objects between new functions
+            nf = reapply_patches_if_need(f)
+            frame_locals[name] = param_as_standalone_func(p, nf, name)
+            frame_locals[name].__doc__ = f.__doc__
+
+        # Delete original patches to prevent new function from evaluating
+        # original patching object as well as re-constrfucted patches.
+        delete_patches_if_need(f)
+
+        f.__test__ = False
+
+    return wrapper
+
+
+def reapply_patches_if_need(func):
+    def dummy_wrapper(orgfunc):
+        @functools.wraps(orgfunc)
+        def dummy_func(*args, **kwargs):
+            return orgfunc(*args, **kwargs)
+
+        return dummy_func
+
+    if hasattr(func, 'patchings'):
+        func = dummy_wrapper(func)
+        tmp_patchings = func.patchings
+        delattr(func, 'patchings')
+        for patch_obj in tmp_patchings:
+            func = patch_obj.decorate_callable(func)
+    return func
+
+
+def delete_patches_if_need(func):
+    if hasattr(func, 'patchings'):
+        func.patchings[:] = []
+
+
+def default_name_func(func, num, p):
+    base_name = func.__name__
+    name_suffix = f"_{num}"
+
+    if len(p.args) > 0 and isinstance(p.args[0], str):
+        name_suffix += "_" + to_safe_name(p.args[0])
+    return base_name + name_suffix
+
+
+def param_as_standalone_func(p, func, name):
+    @functools.wraps(func)
+    def standalone_func(*a):
+        return func(*(a + p.args), **p.kwargs)
+
+    standalone_func.__name__ = name
+
+    # place_as is used by py.test to determine what source file should be
+    # used for this test.
+    standalone_func.place_as = func
+
+    # Remove __wrapped__ because py.test will try to look at __wrapped__
+    # to determine which parameters should be used with this test case,
+    # and obviously we don't need it to do any parameterization.
+    try:
+        del standalone_func.__wrapped__
+    except AttributeError:
+        pass
+    return standalone_func
+
+
+def input_as_callable(input):
+    if callable(input):
+        return lambda: check_input_values(input())
+    input_values = check_input_values(input)
+    return lambda: input_values
+
+
+def check_input_values(input_values):
+    if not isinstance(input_values, list):
+        input_values = list(input_values)
+    return [param.from_decorator(p) for p in input_values]
+
+
+def skip_on_empty_helper(*a, **kw):
+    raise SkipTest("parameterized input is empty")
+
+
+_param = collections.namedtuple("param", "args kwargs")
+
+
+class param(_param):
+    def __new__(cls, *args, **kwargs):
+        return _param.__new__(cls, args, kwargs)
+
+    @classmethod
+    def explicit(cls, args=None, kwargs=None):
+        """Creates a ``param`` by explicitly specifying ``args`` and
+        ``kwargs``::
+            >>> param.explicit([1,2,3])
+            param(*(1, 2, 3))
+            >>> param.explicit(kwargs={"foo": 42})
+            param(*(), **{"foo": "42"})
+        """
+        args = args or ()
+        kwargs = kwargs or {}
+        return cls(*args, **kwargs)
+
+    @classmethod
+    def from_decorator(cls, args):
+        """Returns an instance of ``param()`` for ``@parameterized`` argument
+        ``args``::
+            >>> param.from_decorator((42, ))
+            param(args=(42, ), kwargs={})
+            >>> param.from_decorator("foo")
+            param(args=("foo", ), kwargs={})
+        """
+        if isinstance(args, param):
+            return args
+        elif isinstance(args, str):
+            args = (args,)
+        try:
+            return cls(*args)
+        except TypeError as e:
+            if "after * must be" not in str(e):
+                raise
+            raise TypeError(
+                f"Parameters must be tuples, but {args!r} is not (hint: use '({args!r}, )')",
+            )
+
+    def __repr__(self):
+        return "param(*{!r}, **{!r})".format(*self)
+
+
+def to_safe_name(s):
+    return str(re.sub("[^a-zA-Z0-9_]+", "_", s))
+
+
+# alias
+parameterize = parameterize_func
+param_cls = parameterize_cls
+param_func = parameterize_func

test/distribution/test_dirichlet_op.py

@@ -0,0 +1,142 @@
+# Copyright (c) 2021 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
+#
+#     http://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
+import scipy.stats
+from op_test import OpTest, convert_float_to_uint16, convert_uint16_to_float
+
+import paddle
+from paddle.base import core
+
+paddle.enable_static()
+
+
+class TestDirichletOp(OpTest):
+    # Because dirichlet random sample have not gradient, we skip gradient check.
+    no_need_check_grad = True
+
+    def setUp(self):
+        self.op_type = "dirichlet"
+        self.alpha = np.array((1.0, 2.0))
+        self.sample_shape = (100000, 2)
+
+        self.inputs = {'Alpha': np.broadcast_to(self.alpha, self.sample_shape)}
+        self.attrs = {}
+        self.outputs = {'Out': np.zeros(self.sample_shape)}
+
+    def test_check_output(self):
+        self.check_output_customized(self._hypothesis_testing)
+
+    def _hypothesis_testing(self, outs):
+        self.assertEqual(outs[0].shape, self.sample_shape)
+        self.assertTrue(np.all(outs[0] > 0.0))
+        sample = outs[0][:, 0].astype(np.float64)
+        self.assertLess(
+            scipy.stats.kstest(
+                sample,
+                # scipy dirichlet have not cdf, use beta to replace it.
+                scipy.stats.beta(a=self.alpha[0], b=self.alpha[1]).cdf,
+            )[0],
+            0.01,
+        )
+
+
+class TestDirichletFP16Op(OpTest):
+    # Because dirichlet random sample have not gradient, we skip gradient check.
+    no_need_check_grad = True
+
+    def setUp(self):
+        self.op_type = "dirichlet"
+        self.alpha = np.array((1.0, 2.0))
+        self.sample_shape = (100000, 2)
+        self.dtype = np.float16
+
+        self.inputs = {
+            'Alpha': np.broadcast_to(self.alpha, self.sample_shape).astype(
+                self.dtype
+            )
+        }
+        self.attrs = {}
+        self.outputs = {'Out': np.zeros(self.sample_shape).astype(self.dtype)}
+
+    def test_check_output(self):
+        self.check_output_customized(self._hypothesis_testing)
+
+    def _hypothesis_testing(self, outs):
+        self.assertEqual(outs[0].shape, self.sample_shape)
+        self.assertTrue(np.all(outs[0] > 0.0))
+        sample = outs[0][:, 0].astype(np.float64)
+        self.assertLess(
+            scipy.stats.kstest(
+                sample,
+                # scipy dirichlet have not cdf, use beta to replace it.
+                scipy.stats.beta(a=self.alpha[0], b=self.alpha[1]).cdf,
+            )[0],
+            0.01,
+        )
+
+
+@unittest.skipIf(
+    not core.is_compiled_with_cuda()
+    or not core.is_bfloat16_supported(core.CUDAPlace(0)),
+    "core is not compiled with CUDA and not support the bfloat16",
+)
+class TestDirichletBF16Op(OpTest):
+    # Because dirichlet random sample have not gradient, we skip gradient check.
+    no_need_check_grad = True
+
+    def setUp(self):
+        self.op_type = "dirichlet"
+        self.alpha = np.array((1.0, 2.0))
+        self.sample_shape = (10000, 2)
+        self.dtype = np.uint16
+        self.np_dtype = np.float32
+
+        self.inputs = {
+            'Alpha': np.broadcast_to(self.alpha, self.sample_shape).astype(
+                self.np_dtype
+            )
+        }
+        self.attrs = {}
+        self.outputs = {
+            'Out': np.zeros(self.sample_shape).astype(self.np_dtype)
+        }
+        self.inputs['Alpha'] = convert_float_to_uint16(self.inputs['Alpha'])
+        self.outputs['Out'] = convert_float_to_uint16(self.outputs['Out'])
+        self.place = core.CUDAPlace(0)
+
+    def test_check_output(self):
+        self.check_output_with_place_customized(
+            self._hypothesis_testing, place=core.CUDAPlace(0)
+        )
+
+    def _hypothesis_testing(self, outs):
+        outs = convert_uint16_to_float(outs)
+        self.assertEqual(outs[0].shape, self.sample_shape)
+        self.assertTrue(np.all(outs[0] > 0.0))
+        self.assertLess(
+            scipy.stats.kstest(
+                outs[0][:, 0],
+                # scipy dirichlet have not cdf, use beta to replace it.
+                scipy.stats.beta(a=self.alpha[0], b=self.alpha[1]).cdf,
+            )[0],
+            0.3,  # The bfloat16 test difference is below 0.3
+        )
+
+
+if __name__ == '__main__':
+    unittest.main()

test/distribution/test_distribution.py

@@ -0,0 +1,30 @@
+#   Copyright (c) 2020 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
+#
+#     http://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.
+
+
+class DistributionNumpy:
+    def sample(self):
+        raise NotImplementedError
+
+    def entropy(self):
+        raise NotImplementedError
+
+    def kl_divergence(self, other):
+        raise NotImplementedError
+
+    def log_prob(self, value):
+        raise NotImplementedError
+
+    def probs(self, value):
+        raise NotImplementedError

test/distribution/test_distribution_bernoulli.py

@@ -0,0 +1,584 @@
+# Copyright (c) 2021 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
+#
+#     http://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
+import scipy.special
+import scipy.stats
+from distribution.config import ATOL, DEVICES, RTOL
+from parameterize import (
+    TEST_CASE_NAME,
+    parameterize_cls,
+    parameterize_func,
+    place,
+)
+from test_distribution import DistributionNumpy
+
+import paddle
+from paddle.base.data_feeder import convert_dtype
+from paddle.distribution import Bernoulli
+from paddle.distribution.kl import kl_divergence
+
+np.random.seed(2023)
+paddle.seed(2023)
+
+# Smallest representable number.
+EPS = {
+    'float32': np.finfo('float32').eps,
+    'float64': np.finfo('float64').eps,
+}
+
+
+def _clip_probs_ndarray(probs, dtype):
+    """Clip probs from [0, 1] to (0, 1) with ``eps``"""
+    eps = EPS.get(dtype)
+    return np.clip(probs, a_min=eps, a_max=1 - eps).astype(dtype)
+
+
+def _sigmoid(z):
+    return scipy.special.expit(z)
+
+
+def _kstest(samples_a, samples_b, temperature=1):
+    """Uses the Kolmogorov-Smirnov test for goodness of fit."""
+    _, p_value = scipy.stats.ks_2samp(samples_a, samples_b)
+    return not (p_value < 0.02 * (min(1, temperature)))
+
+
+class BernoulliNumpy(DistributionNumpy):
+    def __init__(self, probs):
+        probs = np.array(probs)
+        if str(probs.dtype) not in ['float32', 'float64']:
+            self.dtype = 'float32'
+        else:
+            self.dtype = probs.dtype
+
+        self.batch_shape = np.shape(probs)
+
+        self.probs = _clip_probs_ndarray(
+            np.array(probs, dtype=self.dtype), str(self.dtype)
+        )
+        self.logits = self._probs_to_logits(self.probs, is_binary=True)
+
+        self.rv = scipy.stats.bernoulli(self.probs.astype('float64'))
+
+    @property
+    def mean(self):
+        return self.rv.mean().astype(self.dtype)
+
+    @property
+    def variance(self):
+        return self.rv.var().astype(self.dtype)
+
+    def sample(self, shape):
+        shape = np.array(shape, dtype='int')
+        if shape.ndim:
+            shape = shape.tolist()
+        else:
+            shape = [shape.tolist()]
+        return self.rv.rvs(size=shape + list(self.batch_shape)).astype(
+            self.dtype
+        )
+
+    def log_prob(self, value):
+        return self.rv.logpmf(value).astype(self.dtype)
+
+    def prob(self, value):
+        return self.rv.pmf(value).astype(self.dtype)
+
+    def cdf(self, value):
+        return self.rv.cdf(value).astype(self.dtype)
+
+    def entropy(self):
+        return (
+            np.maximum(
+                self.logits,
+                0,
+            )
+            - self.logits * self.probs
+            + np.log(1 + np.exp(-np.abs(self.logits)))
+        ).astype(self.dtype)
+
+    def kl_divergence(self, other):
+        """
+        .. math::
+
+            KL[a || b] = Pa * Log[Pa / Pb] + (1 - Pa) * Log[(1 - Pa) / (1 - Pb)]
+        """
+        p_a = self.probs
+        p_b = other.probs
+        return (
+            p_a * np.log(p_a / p_b) + (1 - p_a) * np.log((1 - p_a) / (1 - p_b))
+        ).astype(self.dtype)
+
+    def _probs_to_logits(self, probs, is_binary=False):
+        return (
+            (np.log(probs) - np.log1p(-probs)) if is_binary else np.log(probs)
+        ).astype(self.dtype)
+
+
+class BernoulliTest(unittest.TestCase):
+    def setUp(self):
+        paddle.disable_static(self.place)
+        with paddle.base.dygraph.guard(self.place):
+            # just for convenience
+            self.dtype = self.expected_dtype
+
+            # init numpy with `dtype`
+            self.init_numpy_data(self.probs, self.dtype)
+
+            # init paddle and check dtype convert.
+            self.init_dynamic_data(self.probs, self.default_dtype, self.dtype)
+
+    def init_numpy_data(self, probs, dtype):
+        probs = np.array(probs).astype(dtype)
+        self.rv_np = BernoulliNumpy(probs)
+
+    def init_dynamic_data(self, probs, default_dtype, dtype):
+        self.rv_paddle = Bernoulli(probs)
+        self.assertTrue(
+            dtype == convert_dtype(self.rv_paddle.probs.dtype),
+            (dtype, self.rv_paddle.probs.dtype),
+        )
+
+
+@place(DEVICES)
+@parameterize_cls(
+    (TEST_CASE_NAME, 'probs', 'default_dtype', 'expected_dtype'),
+    [
+        # 0-D probs
+        ('probs_00_32', paddle.full((), 0.0), 'float32', 'float32'),
+        ('probs_03_32', paddle.full((), 0.3), 'float32', 'float32'),
+        ('probs_10_32', paddle.full((), 1.0), 'float32', 'float32'),
+        (
+            'probs_00_64',
+            paddle.full((), 0.0, dtype='float64'),
+            'float64',
+            'float64',
+        ),
+        (
+            'probs_03_64',
+            paddle.full((), 0.3, dtype='float64'),
+            'float64',
+            'float64',
+        ),
+        (
+            'probs_10_64',
+            paddle.full((), 1.0, dtype='float64'),
+            'float64',
+            'float64',
+        ),
+        # 1-D probs
+        ('probs_00', 0.0, 'float64', 'float32'),
+        ('probs_03', 0.3, 'float64', 'float32'),
+        ('probs_10', 1.0, 'float64', 'float32'),
+        ('probs_tensor_03_32', paddle.to_tensor([0.3]), 'float32', 'float32'),
+        (
+            'probs_tensor_03_64',
+            paddle.to_tensor([0.3], dtype='float64'),
+            'float64',
+            'float64',
+        ),
+        (
+            'probs_tensor_03_list_32',
+            paddle.to_tensor(
+                [
+                    0.3,
+                ]
+            ),
+            'float32',
+            'float32',
+        ),
+        (
+            'probs_tensor_03_list_64',
+            paddle.to_tensor(
+                [
+                    0.3,
+                ],
+                dtype='float64',
+            ),
+            'float64',
+            'float64',
+        ),
+        # N-D probs
+        (
+            'probs_tensor_0305',
+            paddle.to_tensor((0.3, 0.5)),
+            'float32',
+            'float32',
+        ),
+        (
+            'probs_tensor_03050104',
+            paddle.to_tensor(((0.3, 0.5), (0.1, 0.4))),
+            'float32',
+            'float32',
+        ),
+    ],
+)
+class BernoulliTestFeature(BernoulliTest):
+    def test_mean(self):
+        with paddle.base.dygraph.guard(self.place):
+            np.testing.assert_allclose(
+                self.rv_paddle.mean,
+                self.rv_np.mean,
+                rtol=RTOL.get(self.dtype),
+                atol=ATOL.get(self.dtype),
+            )
+
+    def test_variance(self):
+        with paddle.base.dygraph.guard(self.place):
+            np.testing.assert_allclose(
+                self.rv_paddle.variance,
+                self.rv_np.variance,
+                rtol=RTOL.get(self.dtype),
+                atol=ATOL.get(self.dtype),
+            )
+
+    @parameterize_func(
+        [
+            (
+                paddle.to_tensor(
+                    [
+                        0.0,
+                    ]
+                ),
+            ),
+            (
+                paddle.to_tensor(
+                    [0.0],
+                ),
+            ),
+            (paddle.to_tensor([1.0]),),
+            (paddle.to_tensor([0.0], dtype='float64'),),
+        ]
+    )
+    def test_log_prob(self, value):
+        with paddle.base.dygraph.guard(self.place):
+            if convert_dtype(value.dtype) == convert_dtype(
+                self.rv_paddle.probs.dtype
+            ):
+                log_prob = self.rv_paddle.log_prob(value)
+                np.testing.assert_allclose(
+                    log_prob,
+                    self.rv_np.log_prob(value),
+                    rtol=RTOL.get(self.dtype),
+                    atol=ATOL.get(self.dtype),
+                )
+                self.assertTrue(self.dtype == convert_dtype(log_prob.dtype))
+
+            else:
+                with self.assertWarns(UserWarning):
+                    self.rv_paddle.log_prob(value)
+
+    @parameterize_func(
+        [
+            (
+                paddle.to_tensor(
+                    [
+                        0.0,
+                    ]
+                ),
+            ),
+            (paddle.to_tensor([0.0]),),
+            (paddle.to_tensor([1.0]),),
+            (paddle.to_tensor([0.0], dtype='float64'),),
+        ]
+    )
+    def test_prob(self, value):
+        with paddle.base.dygraph.guard(self.place):
+            if convert_dtype(value.dtype) == convert_dtype(
+                self.rv_paddle.probs.dtype
+            ):
+                prob = self.rv_paddle.prob(value)
+                np.testing.assert_allclose(
+                    prob,
+                    self.rv_np.prob(value),
+                    rtol=RTOL.get(self.dtype),
+                    atol=ATOL.get(self.dtype),
+                )
+                self.assertTrue(self.dtype == convert_dtype(prob.dtype))
+
+            else:
+                with self.assertWarns(UserWarning):
+                    self.rv_paddle.prob(value)
+
+    @parameterize_func(
+        [
+            (
+                paddle.to_tensor(
+                    [
+                        0.0,
+                    ]
+                ),
+            ),
+            (paddle.to_tensor([0.0]),),
+            (paddle.to_tensor([0.3]),),
+            (paddle.to_tensor([0.7]),),
+            (paddle.to_tensor([1.0]),),
+            (paddle.to_tensor([0.0], dtype='float64'),),
+        ]
+    )
+    def test_cdf(self, value):
+        with paddle.base.dygraph.guard(self.place):
+            if convert_dtype(value.dtype) == convert_dtype(
+                self.rv_paddle.probs.dtype
+            ):
+                cdf = self.rv_paddle.cdf(value)
+                np.testing.assert_allclose(
+                    cdf,
+                    self.rv_np.cdf(value),
+                    rtol=RTOL.get(self.dtype),
+                    atol=ATOL.get(self.dtype),
+                )
+                self.assertTrue(self.dtype == convert_dtype(cdf.dtype))
+
+            else:
+                with self.assertWarns(UserWarning):
+                    self.rv_paddle.cdf(value)
+
+    def test_entropy(self):
+        with paddle.base.dygraph.guard(self.place):
+            np.testing.assert_allclose(
+                self.rv_paddle.entropy(),
+                self.rv_np.entropy(),
+                rtol=RTOL.get(self.dtype),
+                atol=ATOL.get(self.dtype),
+            )
+
+    def test_kl_divergence(self):
+        with paddle.base.dygraph.guard(self.place):
+            other_probs = paddle.to_tensor([0.9], dtype=self.dtype)
+
+            rv_paddle_other = Bernoulli(other_probs)
+            rv_np_other = BernoulliNumpy(other_probs)
+
+            np.testing.assert_allclose(
+                self.rv_paddle.kl_divergence(rv_paddle_other),
+                self.rv_np.kl_divergence(rv_np_other),
+                rtol=RTOL.get(self.dtype),
+                atol=ATOL.get(self.dtype),
+            )
+
+            np.testing.assert_allclose(
+                kl_divergence(self.rv_paddle, rv_paddle_other),
+                self.rv_np.kl_divergence(rv_np_other),
+                rtol=RTOL.get(self.dtype),
+                atol=ATOL.get(self.dtype),
+            )
+
+
+@place(DEVICES)
+@parameterize_cls(
+    (
+        TEST_CASE_NAME,
+        'probs',
+        'default_dtype',
+        'expected_dtype',
+        'shape',
+        'expected_shape',
+    ),
+    [
+        # 0-D probs
+        (
+            'probs_0d_1d',
+            paddle.full((), 0.3),
+            'float32',
+            'float32',
+            [
+                100,
+            ],
+            [
+                100,
+            ],
+        ),
+        (
+            'probs_0d_2d',
+            paddle.full((), 0.3),
+            'float32',
+            'float32',
+            [100, 1],
+            [100, 1],
+        ),
+        (
+            'probs_0d_3d',
+            paddle.full((), 0.3),
+            'float32',
+            'float32',
+            [100, 2, 3],
+            [100, 2, 3],
+        ),
+        # 1-D probs
+        (
+            'probs_1d_1d_32',
+            paddle.to_tensor([0.3]),
+            'float32',
+            'float32',
+            [
+                100,
+            ],
+            [100, 1],
+        ),
+        (
+            'probs_1d_1d_64',
+            paddle.to_tensor([0.3], dtype='float64'),
+            'float64',
+            'float64',
+            paddle.to_tensor(
+                [
+                    100,
+                ]
+            ),
+            [100, 1],
+        ),
+        (
+            'probs_1d_2d',
+            paddle.to_tensor([0.3]),
+            'float32',
+            'float32',
+            [100, 2],
+            [100, 2, 1],
+        ),
+        (
+            'probs_1d_3d',
+            paddle.to_tensor([0.3]),
+            'float32',
+            'float32',
+            [100, 2, 3],
+            [100, 2, 3, 1],
+        ),
+        # N-D probs
+        (
+            'probs_2d_1d',
+            paddle.to_tensor((0.3, 0.5)),
+            'float32',
+            'float32',
+            [
+                100,
+            ],
+            [100, 2],
+        ),
+        (
+            'probs_2d_2d',
+            paddle.to_tensor((0.3, 0.5)),
+            'float32',
+            'float32',
+            [100, 3],
+            [100, 3, 2],
+        ),
+        (
+            'probs_2d_3d',
+            paddle.to_tensor((0.3, 0.5)),
+            'float32',
+            'float32',
+            [100, 4, 3],
+            [100, 4, 3, 2],
+        ),
+    ],
+)
+class BernoulliTestSample(BernoulliTest):
+    def test_sample(self):
+        with paddle.base.dygraph.guard(self.place):
+            sample_np = self.rv_np.sample(self.shape)
+            sample_paddle = self.rv_paddle.sample(self.shape)
+
+            self.assertEqual(list(sample_paddle.shape), self.expected_shape)
+            self.assertEqual(sample_paddle.dtype, self.rv_paddle.probs.dtype)
+
+            if self.probs.ndim:
+                for i in range(len(self.probs)):
+                    self.assertTrue(
+                        _kstest(
+                            sample_np[..., i].reshape(-1),
+                            sample_paddle.numpy()[..., i].reshape(-1),
+                        )
+                    )
+            else:
+                self.assertTrue(
+                    _kstest(
+                        sample_np.reshape(-1),
+                        sample_paddle.numpy().reshape(-1),
+                    )
+                )
+
+    @parameterize_func(
+        [
+            (1.0,),
+            (0.1,),
+        ]
+    )
+    def test_rsample(self, temperature):
+        """Compare two samples from `rsample` method, one from scipy `sample` and another from paddle `rsample`."""
+        with paddle.base.dygraph.guard(self.place):
+            sample_np = self.rv_np.sample(self.shape)
+            rsample_paddle = self.rv_paddle.rsample(self.shape, temperature)
+
+            self.assertEqual(list(rsample_paddle.shape), self.expected_shape)
+            self.assertEqual(rsample_paddle.dtype, self.rv_paddle.probs.dtype)
+
+            if self.probs.ndim:
+                for i in range(len(self.probs)):
+                    self.assertTrue(
+                        _kstest(
+                            sample_np[..., i].reshape(-1),
+                            (
+                                _sigmoid(rsample_paddle.numpy()[..., i]) > 0.5
+                            ).reshape(-1),
+                            temperature,
+                        )
+                    )
+            else:
+                self.assertTrue(
+                    _kstest(
+                        sample_np.reshape(-1),
+                        (_sigmoid(rsample_paddle.numpy()) > 0.5).reshape(-1),
+                        temperature,
+                    )
+                )
+
+    def test_rsample_backpropagation(self):
+        with paddle.base.dygraph.guard(self.place):
+            self.rv_paddle.probs.stop_gradient = False
+            rsample_paddle = self.rv_paddle.rsample(self.shape)
+            rsample_paddle = paddle.nn.functional.sigmoid(rsample_paddle)
+            grads = paddle.grad([rsample_paddle], [self.rv_paddle.probs])
+            self.assertEqual(len(grads), 1)
+            self.assertEqual(grads[0].dtype, self.rv_paddle.probs.dtype)
+            self.assertEqual(grads[0].shape, self.rv_paddle.probs.shape)
+
+
+@place(DEVICES)
+@parameterize_cls([TEST_CASE_NAME], ['BernoulliTestError'])
+class BernoulliTestError(unittest.TestCase):
+    def setUp(self):
+        paddle.disable_static(self.place)
+
+    @parameterize_func(
+        [
+            (
+                [0.3, 0.5],
+                paddle.to_tensor([0.1, 0.2, 0.3]),
+            ),
+        ]
+    )
+    def test_bad_broadcast(self, probs, value):
+        with paddle.base.dygraph.guard(self.place):
+            rv = Bernoulli(probs)
+            self.assertRaises(ValueError, rv.cdf, value)
+            self.assertRaises(ValueError, rv.log_prob, value)
+            self.assertRaises(ValueError, rv.prob, value)
+
+
+if __name__ == '__main__':
+    unittest.main()

test/distribution/test_distribution_bernoulli_static.py

@@ -0,0 +1,473 @@
+# Copyright (c) 2021 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
+#
+#     http://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 sys
+import unittest
+
+import numpy as np
+from distribution.config import ATOL, DEVICES, RTOL
+from parameterize import (
+    TEST_CASE_NAME,
+    parameterize_cls,
+    parameterize_func,
+    place,
+)
+
+sys.path.append("../../distribution")
+from test_distribution_bernoulli import BernoulliNumpy, _kstest, _sigmoid
+
+import paddle
+from paddle.distribution import Bernoulli
+from paddle.distribution.kl import kl_divergence
+
+np.random.seed(2023)
+paddle.seed(2023)
+paddle.enable_static()
+default_dtype = paddle.get_default_dtype()
+
+
+@place(DEVICES)
+@parameterize_cls(
+    (TEST_CASE_NAME, 'params'),  # params: name, probs, probs_other, value
+    [
+        (
+            'params',
+            (
+                # 1-D probs
+                (
+                    'probs_not_iterable',
+                    0.3,
+                    0.7,
+                    1.0,
+                ),
+                (
+                    'probs_not_iterable_and_broadcast_for_value',
+                    0.3,
+                    0.7,
+                    np.array([[0.0, 1.0], [1.0, 0.0]], dtype=default_dtype),
+                ),
+                # N-D probs
+                (
+                    'probs_tuple_0305',
+                    (0.3, 0.5),
+                    0.7,
+                    1.0,
+                ),
+                (
+                    'probs_tuple_03050104',
+                    ((0.3, 0.5), (0.1, 0.4)),
+                    0.7,
+                    1.0,
+                ),
+            ),
+        )
+    ],
+)
+class BernoulliTestFeature(unittest.TestCase):
+    def setUp(self):
+        self.program = paddle.static.Program()
+        self.executor = paddle.static.Executor(self.place)
+
+        self.params_len = len(self.params)
+
+        with paddle.static.program_guard(self.program):
+            self.init_numpy_data(self.params)
+            self.init_static_data(self.params)
+
+    def init_numpy_data(self, params):
+        self.mean_np = []
+        self.variance_np = []
+        self.log_prob_np = []
+        self.prob_np = []
+        self.cdf_np = []
+        self.entropy_np = []
+        self.kl_np = []
+
+        for _, probs, probs_other, value in params:
+            rv_np = BernoulliNumpy(probs)
+            rv_np_other = BernoulliNumpy(probs_other)
+
+            self.mean_np.append(rv_np.mean)
+            self.variance_np.append(rv_np.variance)
+            self.log_prob_np.append(rv_np.log_prob(value))
+            self.prob_np.append(rv_np.prob(value))
+            self.cdf_np.append(rv_np.cdf(value))
+            self.entropy_np.append(rv_np.entropy())
+            self.kl_np.append(rv_np.kl_divergence(rv_np_other))
+
+    def init_static_data(self, params):
+        with paddle.static.program_guard(self.program):
+            rv_paddles = []
+            rv_paddles_other = []
+            values = []
+            for _, probs, probs_other, value in params:
+                if not isinstance(value, np.ndarray):
+                    value = paddle.full([1], value, dtype=default_dtype)
+                else:
+                    value = paddle.to_tensor(value, place=self.place)
+
+                rv_paddles.append(Bernoulli(probs=paddle.to_tensor(probs)))
+                rv_paddles_other.append(
+                    Bernoulli(probs=paddle.to_tensor(probs_other))
+                )
+                values.append(value)
+
+            results = self.executor.run(
+                self.program,
+                feed={},
+                fetch_list=[
+                    [
+                        rv_paddles[i].mean,
+                        rv_paddles[i].variance,
+                        rv_paddles[i].log_prob(values[i]),
+                        rv_paddles[i].prob(values[i]),
+                        rv_paddles[i].cdf(values[i]),
+                        rv_paddles[i].entropy(),
+                        rv_paddles[i].kl_divergence(rv_paddles_other[i]),
+                        kl_divergence(rv_paddles[i], rv_paddles_other[i]),
+                    ]
+                    for i in range(self.params_len)
+                ],
+            )
+
+            self.mean_paddle = []
+            self.variance_paddle = []
+            self.log_prob_paddle = []
+            self.prob_paddle = []
+            self.cdf_paddle = []
+            self.entropy_paddle = []
+            self.kl_paddle = []
+            self.kl_func_paddle = []
+            for i in range(self.params_len):
+                (
+                    _mean,
+                    _variance,
+                    _log_prob,
+                    _prob,
+                    _cdf,
+                    _entropy,
+                    _kl,
+                    _kl_func,
+                ) = results[i * 8 : (i + 1) * 8]
+                self.mean_paddle.append(_mean)
+                self.variance_paddle.append(_variance)
+                self.log_prob_paddle.append(_log_prob)
+                self.prob_paddle.append(_prob)
+                self.cdf_paddle.append(_cdf)
+                self.entropy_paddle.append(_entropy)
+                self.kl_paddle.append(_kl)
+                self.kl_func_paddle.append(_kl_func)
+
+    def test_all(self):
+        for i in range(self.params_len):
+            self._test_mean(i)
+            self._test_variance(i)
+            self._test_log_prob(i)
+            self._test_prob(i)
+            self._test_cdf(i)
+            self._test_entropy(i)
+            self._test_kl_divergence(i)
+
+    def _test_mean(self, i):
+        np.testing.assert_allclose(
+            self.mean_np[i],
+            self.mean_paddle[i],
+            rtol=RTOL.get(default_dtype),
+            atol=ATOL.get(default_dtype),
+        )
+
+    def _test_variance(self, i):
+        np.testing.assert_allclose(
+            self.variance_np[i],
+            self.variance_paddle[i],
+            rtol=RTOL.get(default_dtype),
+            atol=ATOL.get(default_dtype),
+        )
+
+    def _test_log_prob(self, i):
+        np.testing.assert_allclose(
+            self.log_prob_np[i],
+            self.log_prob_paddle[i],
+            rtol=RTOL.get(default_dtype),
+            atol=ATOL.get(default_dtype),
+        )
+
+    def _test_prob(self, i):
+        np.testing.assert_allclose(
+            self.prob_np[i],
+            self.prob_paddle[i],
+            rtol=RTOL.get(default_dtype),
+            atol=ATOL.get(default_dtype),
+        )
+
+    def _test_cdf(self, i):
+        np.testing.assert_allclose(
+            self.cdf_np[i],
+            self.cdf_paddle[i],
+            rtol=RTOL.get(default_dtype),
+            atol=ATOL.get(default_dtype),
+        )
+
+    def _test_entropy(self, i):
+        np.testing.assert_allclose(
+            self.entropy_np[i],
+            self.entropy_paddle[i],
+            rtol=RTOL.get(default_dtype),
+            atol=ATOL.get(default_dtype),
+        )
+
+    def _test_kl_divergence(self, i):
+        np.testing.assert_allclose(
+            self.kl_np[i],
+            self.kl_paddle[i],
+            rtol=RTOL.get(default_dtype),
+            atol=ATOL.get(default_dtype),
+        )
+
+        np.testing.assert_allclose(
+            self.kl_np[i],
+            self.kl_func_paddle[i],
+            rtol=RTOL.get(default_dtype),
+            atol=ATOL.get(default_dtype),
+        )
+
+
+@place(DEVICES)
+@parameterize_cls(
+    (TEST_CASE_NAME, 'probs', 'shape', 'temperature', 'expected_shape'),
+    [
+        # 1-D probs
+        (
+            'probs_03',
+            (0.3,),
+            [
+                100,
+            ],
+            0.1,
+            [100, 1],
+        ),
+        # N-D probs
+        (
+            'probs_0305',
+            (0.3, 0.5),
+            [
+                100,
+            ],
+            0.1,
+            [100, 2],
+        ),
+    ],
+)
+class BernoulliTestSample(unittest.TestCase):
+    def setUp(self):
+        self.program = paddle.static.Program()
+        self.executor = paddle.static.Executor(self.place)
+
+        with paddle.static.program_guard(self.program):
+            self.init_numpy_data(self.probs, self.shape)
+            self.init_static_data(self.probs, self.shape, self.temperature)
+
+    def init_numpy_data(self, probs, shape):
+        self.rv_np = BernoulliNumpy(probs)
+        self.sample_np = self.rv_np.sample(shape)
+
+    def init_static_data(self, probs, shape, temperature):
+        with paddle.static.program_guard(self.program):
+            self.rv_paddle = Bernoulli(probs=paddle.to_tensor(probs))
+
+            [self.sample_paddle, self.rsample_paddle] = self.executor.run(
+                self.program,
+                feed={},
+                fetch_list=[
+                    self.rv_paddle.sample(shape),
+                    self.rv_paddle.rsample(shape, temperature),
+                ],
+            )
+
+    def test_sample(self):
+        with paddle.static.program_guard(self.program):
+            self.assertEqual(
+                list(self.sample_paddle.shape), self.expected_shape
+            )
+
+            for i in range(len(self.probs)):
+                self.assertTrue(
+                    _kstest(
+                        self.sample_np[..., i].reshape(-1),
+                        self.sample_paddle[..., i].reshape(-1),
+                    )
+                )
+
+    def test_rsample(self):
+        """Compare two samples from `rsample` method, one from scipy and another from paddle."""
+        with paddle.static.program_guard(self.program):
+            self.assertEqual(
+                list(self.rsample_paddle.shape), self.expected_shape
+            )
+
+            for i in range(len(self.probs)):
+                self.assertTrue(
+                    _kstest(
+                        self.sample_np[..., i].reshape(-1),
+                        (_sigmoid(self.rsample_paddle[..., i]) > 0.5).reshape(
+                            -1
+                        ),
+                        self.temperature,
+                    )
+                )
+
+
+@place(DEVICES)
+@parameterize_cls([TEST_CASE_NAME], ['BernoulliTestError'])
+class BernoulliTestError(unittest.TestCase):
+    def setUp(self):
+        self.program = paddle.static.Program()
+        self.executor = paddle.static.Executor(self.place)
+
+    @parameterize_func(
+        [
+            (0,),  # int
+            ((0.3,),),  # tuple
+            (
+                [
+                    0.3,
+                ],
+            ),  # list
+            (
+                np.array(
+                    [
+                        0.3,
+                    ]
+                ),
+            ),  # ndarray
+            (-1j + 1,),  # complex
+            ('0',),  # str
+        ]
+    )
+    def test_bad_init_type(self, probs):
+        with (
+            paddle.static.program_guard(self.program),
+            self.assertRaises(TypeError),
+        ):
+            [_] = self.executor.run(
+                self.program, feed={}, fetch_list=[Bernoulli(probs=probs)]
+            )
+
+    @parameterize_func(
+        [
+            (100,),  # int
+            (100.0,),  # float
+        ]
+    )
+    def test_bad_sample_shape_type(self, shape):
+        with paddle.static.program_guard(self.program):
+            rv = Bernoulli(0.3)
+
+            with self.assertRaises(TypeError):
+                [_] = self.executor.run(
+                    self.program, feed={}, fetch_list=[rv.sample(shape)]
+                )
+
+            with self.assertRaises(TypeError):
+                [_] = self.executor.run(
+                    self.program, feed={}, fetch_list=[rv.rsample(shape)]
+                )
+
+    @parameterize_func(
+        [
+            (1,),  # int
+        ]
+    )
+    def test_bad_rsample_temperature_type(self, temperature):
+        with paddle.static.program_guard(self.program):
+            rv = Bernoulli(0.3)
+
+            with self.assertRaises(TypeError):
+                [_] = self.executor.run(
+                    self.program,
+                    feed={},
+                    fetch_list=[rv.rsample([100], temperature)],
+                )
+
+    @parameterize_func(
+        [
+            (1,),  # int
+            (1.0,),  # float
+            ([1.0],),  # list
+            ((1.0),),  # tuple
+            (np.array(1.0),),  # ndarray
+        ]
+    )
+    def test_bad_value_type(self, value):
+        with paddle.static.program_guard(self.program):
+            rv = Bernoulli(0.3)
+
+            with self.assertRaises(TypeError):
+                [_] = self.executor.run(
+                    self.program, feed={}, fetch_list=[rv.log_prob(value)]
+                )
+
+            with self.assertRaises(TypeError):
+                [_] = self.executor.run(
+                    self.program, feed={}, fetch_list=[rv.prob(value)]
+                )
+
+            with self.assertRaises(TypeError):
+                [_] = self.executor.run(
+                    self.program, feed={}, fetch_list=[rv.cdf(value)]
+                )
+
+    @parameterize_func(
+        [
+            (np.array(1.0),),  # ndarray or other distribution
+        ]
+    )
+    def test_bad_kl_other_type(self, other):
+        with paddle.static.program_guard(self.program):
+            rv = Bernoulli(0.3)
+
+            with self.assertRaises(TypeError):
+                [_] = self.executor.run(
+                    self.program, feed={}, fetch_list=[rv.kl_divergence(other)]
+                )
+
+    @parameterize_func(
+        [
+            (paddle.to_tensor([0.1, 0.2, 0.3]),),
+        ]
+    )
+    def test_bad_broadcast(self, value):
+        with paddle.static.program_guard(self.program):
+            rv = Bernoulli(paddle.to_tensor([0.3, 0.5]))
+
+            # `logits, value = paddle.broadcast_tensors([self.logits, value])`
+            # raise ValueError in dygraph, raise TypeError in static.
+            with self.assertRaises((TypeError, ValueError)):
+                [_] = self.executor.run(
+                    self.program, feed={}, fetch_list=[rv.cdf(value)]
+                )
+
+            with self.assertRaises((TypeError, ValueError)):
+                [_] = self.executor.run(
+                    self.program, feed={}, fetch_list=[rv.log_prob(value)]
+                )
+
+            with self.assertRaises((TypeError, ValueError)):
+                [_] = self.executor.run(
+                    self.program, feed={}, fetch_list=[rv.prob(value)]
+                )
+
+
+if __name__ == '__main__':
+    unittest.main()

test/distribution/test_distribution_beta.py

@@ -0,0 +1,127 @@
+# Copyright (c) 2021 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
+#
+#     http://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 numbers
+import unittest
+
+import numpy as np
+import scipy.stats
+from distribution.config import ATOL, DEVICES, RTOL
+from parameterize import TEST_CASE_NAME, parameterize_cls, place, xrand
+
+import paddle
+
+np.random.seed(2022)
+
+
+@place(DEVICES)
+@parameterize_cls(
+    (TEST_CASE_NAME, 'alpha', 'beta'),
+    [
+        ('test-scale', 1.0, 2.0),
+        ('test-tensor', xrand(), xrand()),
+        ('test-broadcast', xrand((2, 1)), xrand((2, 5))),
+    ],
+)
+class TestBeta(unittest.TestCase):
+    def setUp(self):
+        # scale no need convert to tensor for scale input unittest
+        alpha, beta = self.alpha, self.beta
+        if not isinstance(self.alpha, numbers.Real):
+            alpha = paddle.to_tensor(self.alpha)
+        if not isinstance(self.beta, numbers.Real):
+            beta = paddle.to_tensor(self.beta)
+
+        self._paddle_beta = paddle.distribution.Beta(alpha, beta)
+
+    def test_mean(self):
+        with paddle.base.dygraph.guard(self.place):
+            np.testing.assert_allclose(
+                self._paddle_beta.mean,
+                scipy.stats.beta.mean(self.alpha, self.beta),
+                rtol=RTOL.get(str(self._paddle_beta.alpha.numpy().dtype)),
+                atol=ATOL.get(str(self._paddle_beta.alpha.numpy().dtype)),
+            )
+
+    def test_variance(self):
+        with paddle.base.dygraph.guard(self.place):
+            np.testing.assert_allclose(
+                self._paddle_beta.variance,
+                scipy.stats.beta.var(self.alpha, self.beta),
+                rtol=RTOL.get(str(self._paddle_beta.alpha.numpy().dtype)),
+                atol=ATOL.get(str(self._paddle_beta.alpha.numpy().dtype)),
+            )
+
+    def test_prob(self):
+        value = [np.random.rand(*self._paddle_beta.alpha.shape)]
+
+        for v in value:
+            with paddle.base.dygraph.guard(self.place):
+                np.testing.assert_allclose(
+                    self._paddle_beta.prob(paddle.to_tensor(v)),
+                    scipy.stats.beta.pdf(v, self.alpha, self.beta),
+                    rtol=RTOL.get(str(self._paddle_beta.alpha.numpy().dtype)),
+                    atol=ATOL.get(str(self._paddle_beta.alpha.numpy().dtype)),
+                )
+
+    def test_log_prob(self):
+        value = [np.random.rand(*self._paddle_beta.alpha.shape)]
+
+        for v in value:
+            with paddle.base.dygraph.guard(self.place):
+                np.testing.assert_allclose(
+                    self._paddle_beta.log_prob(paddle.to_tensor(v)),
+                    scipy.stats.beta.logpdf(v, self.alpha, self.beta),
+                    rtol=RTOL.get(str(self._paddle_beta.alpha.numpy().dtype)),
+                    atol=ATOL.get(str(self._paddle_beta.alpha.numpy().dtype)),
+                )
+
+    def test_entropy(self):
+        with paddle.base.dygraph.guard(self.place):
+            np.testing.assert_allclose(
+                self._paddle_beta.entropy(),
+                scipy.stats.beta.entropy(self.alpha, self.beta),
+                rtol=RTOL.get(str(self._paddle_beta.alpha.numpy().dtype)),
+                atol=ATOL.get(str(self._paddle_beta.alpha.numpy().dtype)),
+            )
+
+    def test_sample_shape(self):
+        cases = [
+            {
+                'input': [],
+                'expect': list(paddle.squeeze(self._paddle_beta.alpha).shape),
+            },
+            {
+                'input': [2, 3],
+                'expect': [
+                    2,
+                    3,
+                    *paddle.squeeze(self._paddle_beta.alpha).shape,
+                ],
+            },
+        ]
+        for case in cases:
+            self.assertTrue(
+                self._paddle_beta.sample(case.get('input')).shape
+                == case.get('expect')
+            )
+
+    def test_errors(self):
+        with self.assertRaises(ValueError):
+            array = np.array([], dtype=np.float32)
+            x = paddle.to_tensor(np.reshape(array, [0]), dtype='int32')
+            paddle.distribution.Beta(alpha=x, beta=x)
+
+
+if __name__ == '__main__':
+    unittest.main()

test/distribution/test_distribution_beta_static.py

@@ -0,0 +1,148 @@
+# Copyright (c) 2021 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
+#
+#     http://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
+import parameterize as param
+import scipy.stats
+from distribution import config
+from distribution.config import ATOL, RTOL
+from parameterize import xrand
+
+import paddle
+
+np.random.seed(2022)
+paddle.enable_static()
+
+
+@param.place(config.DEVICES)
+@param.parameterize_cls(
+    (param.TEST_CASE_NAME, 'alpha', 'beta'),
+    [
+        ('test-tensor', xrand((10, 10)), xrand((10, 10))),
+        ('test-broadcast', xrand((2, 1)), xrand((2, 5))),
+        ('test-larger-data', xrand((10, 20)), xrand((10, 20))),
+    ],
+)
+class TestBeta(unittest.TestCase):
+    def setUp(self):
+        self.program = paddle.static.Program()
+        self.executor = paddle.static.Executor(self.place)
+        with paddle.static.program_guard(self.program):
+            # scale no need convert to tensor for scale input unittest
+            alpha = paddle.static.data(
+                'alpha', self.alpha.shape, self.alpha.dtype
+            )
+            beta = paddle.static.data('beta', self.beta.shape, self.beta.dtype)
+            self._paddle_beta = paddle.distribution.Beta(alpha, beta)
+            self.feeds = {'alpha': self.alpha, 'beta': self.beta}
+
+    def test_mean(self):
+        with paddle.static.program_guard(self.program):
+            [mean] = self.executor.run(
+                self.program,
+                feed=self.feeds,
+                fetch_list=[self._paddle_beta.mean],
+            )
+            np.testing.assert_allclose(
+                mean,
+                scipy.stats.beta.mean(self.alpha, self.beta),
+                rtol=RTOL.get(str(self.alpha.dtype)),
+                atol=ATOL.get(str(self.alpha.dtype)),
+            )
+
+    def test_variance(self):
+        with paddle.static.program_guard(self.program):
+            [variance] = self.executor.run(
+                self.program,
+                feed=self.feeds,
+                fetch_list=[self._paddle_beta.variance],
+            )
+            np.testing.assert_allclose(
+                variance,
+                scipy.stats.beta.var(self.alpha, self.beta),
+                rtol=RTOL.get(str(self.alpha.dtype)),
+                atol=ATOL.get(str(self.alpha.dtype)),
+            )
+
+    def test_prob(self):
+        with paddle.static.program_guard(self.program):
+            value = paddle.static.data(
+                'value',
+                self._paddle_beta.alpha.shape,
+                self._paddle_beta.alpha.dtype,
+            )
+            prob = self._paddle_beta.prob(value)
+
+            random_number = np.random.rand(*self._paddle_beta.alpha.shape)
+            feeds = dict(self.feeds, value=random_number)
+            [prob] = self.executor.run(
+                self.program, feed=feeds, fetch_list=[prob]
+            )
+            np.testing.assert_allclose(
+                prob,
+                scipy.stats.beta.pdf(random_number, self.alpha, self.beta),
+                rtol=RTOL.get(str(self.alpha.dtype)),
+                atol=ATOL.get(str(self.alpha.dtype)),
+            )
+
+    def test_log_prob(self):
+        with paddle.static.program_guard(self.program):
+            value = paddle.static.data(
+                'value',
+                self._paddle_beta.alpha.shape,
+                self._paddle_beta.alpha.dtype,
+            )
+            prob = self._paddle_beta.log_prob(value)
+            random_number = np.random.rand(*self._paddle_beta.alpha.shape)
+            feeds = dict(self.feeds, value=random_number)
+            [prob] = self.executor.run(
+                self.program, feed=feeds, fetch_list=[prob]
+            )
+            np.testing.assert_allclose(
+                prob,
+                scipy.stats.beta.logpdf(random_number, self.alpha, self.beta),
+                rtol=RTOL.get(str(self.alpha.dtype)),
+                atol=ATOL.get(str(self.alpha.dtype)),
+            )
+
+    def test_entropy(self):
+        with paddle.static.program_guard(self.program):
+            [entropy] = self.executor.run(
+                self.program,
+                feed=self.feeds,
+                fetch_list=[self._paddle_beta.entropy()],
+            )
+            np.testing.assert_allclose(
+                entropy,
+                scipy.stats.beta.entropy(self.alpha, self.beta),
+                rtol=RTOL.get(str(self.alpha.dtype)),
+                atol=ATOL.get(str(self.alpha.dtype)),
+            )
+
+    def test_sample(self):
+        with paddle.static.program_guard(self.program):
+            [data] = self.executor.run(
+                self.program,
+                feed=self.feeds,
+                fetch_list=self._paddle_beta.sample(),
+            )
+            self.assertTrue(
+                data.shape
+                == np.broadcast_arrays(self.alpha, self.beta)[0].shape
+            )
+
+
+if __name__ == '__main__':
+    unittest.main()

test/distribution/test_distribution_binomial.py

@@ -0,0 +1,213 @@
+# Copyright (c) 2021 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
+#
+#     http://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
+import parameterize
+import scipy.stats
+from distribution import config
+
+import paddle
+from paddle.distribution.binomial import Binomial
+
+
+@parameterize.place(config.DEVICES)
+@parameterize.parameterize_cls(
+    (parameterize.TEST_CASE_NAME, 'total_count', 'probs'),
+    [
+        (
+            'zero-dim',
+            np.array(1000),
+            np.array(0.6),
+        ),
+        (
+            'one-dim',
+            1000,
+            np.array([0.4]).astype('float32'),
+        ),
+        (
+            'multi-dim-total_count-probability',
+            parameterize.xrand((2, 1), min=1, max=100).astype('int32'),
+            parameterize.xrand((2, 3), dtype='float64', min=0.3, max=1),
+        ),
+    ],
+)
+class TestBinomial(unittest.TestCase):
+    def setUp(self):
+        self._dist = Binomial(
+            total_count=paddle.to_tensor(self.total_count),
+            probs=paddle.to_tensor(self.probs),
+        )
+
+    def test_mean(self):
+        mean = self._dist.mean
+        self.assertEqual(mean.numpy().dtype, self.probs.dtype)
+        np.testing.assert_allclose(
+            mean,
+            self._np_mean(),
+            rtol=config.RTOL.get(str(self.probs.dtype)),
+            atol=config.ATOL.get(str(self.probs.dtype)),
+        )
+
+    def test_variance(self):
+        var = self._dist.variance
+        self.assertEqual(var.numpy().dtype, self.probs.dtype)
+        np.testing.assert_allclose(
+            var,
+            self._np_variance(),
+            rtol=config.RTOL.get(str(self.probs.dtype)),
+            atol=config.ATOL.get(str(self.probs.dtype)),
+        )
+
+    def test_entropy(self):
+        entropy = self._dist.entropy()
+        self.assertEqual(entropy.numpy().dtype, self.probs.dtype)
+        np.testing.assert_allclose(
+            entropy,
+            self._np_entropy(),
+            rtol=config.RTOL.get(str(self.probs.dtype)),
+            atol=config.ATOL.get(str(self.probs.dtype)),
+        )
+
+    def test_sample(self):
+        sample_shape = ()
+        samples = self._dist.sample(sample_shape)
+        self.assertEqual(
+            tuple(samples.shape),
+            sample_shape + self._dist.batch_shape + self._dist.event_shape,
+        )
+
+        sample_shape = (5000,)
+        samples = self._dist.sample(sample_shape)
+        sample_mean = samples.mean(axis=0)
+        sample_variance = samples.var(axis=0)
+
+        np.testing.assert_allclose(
+            sample_mean, self._dist.mean, atol=0, rtol=0.20
+        )
+        np.testing.assert_allclose(
+            sample_variance, self._dist.variance, atol=0, rtol=0.20
+        )
+
+    def _np_variance(self):
+        return scipy.stats.binom.var(self.total_count, self.probs)
+
+    def _np_mean(self):
+        return scipy.stats.binom.mean(self.total_count, self.probs)
+
+    def _np_entropy(self):
+        return scipy.stats.binom.entropy(self.total_count, self.probs)
+
+
+@parameterize.place(config.DEVICES)
+@parameterize.parameterize_cls(
+    (parameterize.TEST_CASE_NAME, 'total_count', 'probs', 'value'),
+    [
+        (
+            'value-same-shape',
+            1000,
+            np.array([0.12, 0.3, 0.85]).astype('float64'),
+            np.array([2.0, 55.0, 999.0]).astype('float64'),
+        ),
+        (
+            'value-broadcast-shape',
+            10,
+            np.array([[0.3, 0.7], [0.5, 0.5]]),
+            np.array([[[4.0, 6], [8, 2]], [[2.0, 4], [9, 7]]]),
+        ),
+    ],
+)
+class TestBinomialProbs(unittest.TestCase):
+    def setUp(self):
+        self._dist = Binomial(
+            total_count=paddle.to_tensor(self.total_count),
+            probs=paddle.to_tensor(self.probs),
+        )
+
+    def test_prob(self):
+        np.testing.assert_allclose(
+            self._dist.prob(paddle.to_tensor(self.value)),
+            scipy.stats.binom.pmf(self.value, self.total_count, self.probs),
+            rtol=config.RTOL.get(str(self.probs.dtype)),
+            atol=config.ATOL.get(str(self.probs.dtype)),
+        )
+
+    def test_log_prob(self):
+        np.testing.assert_allclose(
+            self._dist.log_prob(paddle.to_tensor(self.value)),
+            scipy.stats.binom.logpmf(self.value, self.total_count, self.probs),
+            rtol=config.RTOL.get(str(self.probs.dtype)),
+            atol=config.ATOL.get(str(self.probs.dtype)),
+        )
+
+
+@parameterize.place(config.DEVICES)
+@parameterize.parameterize_cls(
+    (parameterize.TEST_CASE_NAME, 'n_1', 'p_1', 'n_2', 'p_2'),
+    [
+        (
+            'one-dim-probability',
+            np.array([3333]),
+            parameterize.xrand((1,), dtype='float32', min=0, max=1),
+            np.array([3333]),
+            parameterize.xrand((1,), dtype='float32', min=0, max=1),
+        ),
+        (
+            'multi-dim-probability',
+            np.array([25, 25, 25]),
+            parameterize.xrand((2, 3), dtype='float64', min=0, max=1),
+            np.array([25, 25, 25]),
+            parameterize.xrand((2, 3), dtype='float64', min=0, max=1),
+        ),
+    ],
+)
+class TestBinomialKL(unittest.TestCase):
+    def setUp(self):
+        self._dist1 = Binomial(
+            total_count=paddle.to_tensor(self.n_1),
+            probs=paddle.to_tensor(self.p_1),
+        )
+        self._dist2 = Binomial(
+            total_count=paddle.to_tensor(self.n_2),
+            probs=paddle.to_tensor(self.p_2),
+        )
+
+    def test_kl_divergence(self):
+        kl0 = self._dist1.kl_divergence(self._dist2)
+        kl1 = self.kl_divergence(self._dist1, self._dist2)
+
+        self.assertEqual(tuple(kl0.shape), self.p_1.shape)
+        self.assertEqual(tuple(kl1.shape), self.p_1.shape)
+        np.testing.assert_allclose(
+            kl0,
+            kl1,
+            rtol=config.RTOL.get(str(self.p_1.dtype)),
+            atol=config.ATOL.get(str(self.p_1.dtype)),
+        )
+
+    def kl_divergence(self, dist1, dist2):
+        support = np.arange(1 + self.n_1.max(), dtype=self.p_1.dtype)
+        support = support.reshape((-1,) + (1,) * len(self.p_1.shape))
+        log_prob_1 = scipy.stats.binom.logpmf(
+            support, dist1.total_count, dist1.probs
+        )
+        log_prob_2 = scipy.stats.binom.logpmf(
+            support, dist2.total_count, dist2.probs
+        )
+        return (np.exp(log_prob_1) * (log_prob_1 - log_prob_2)).sum(0)
+
+
+if __name__ == '__main__':
+    unittest.main()

test/distribution/test_distribution_binomial_static.py

@@ -0,0 +1,256 @@
+# Copyright (c) 2021 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
+#
+#     http://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
+import parameterize
+import scipy.stats
+from distribution import config
+
+import paddle
+from paddle.distribution.binomial import Binomial
+
+paddle.enable_static()
+
+
+paddle.enable_static()
+
+
+@parameterize.place(config.DEVICES)
+@parameterize.parameterize_cls(
+    (parameterize.TEST_CASE_NAME, 'total_count', 'probs'),
+    [
+        (
+            'zero-dim',
+            np.array(1000),
+            np.array(0.6),
+        ),
+        (
+            'one-dim',
+            np.array([1000]),
+            parameterize.xrand((1,), dtype='float32', min=0, max=1),
+        ),
+        (
+            'multi-dim',
+            np.array([100]),
+            parameterize.xrand((1, 3), dtype='float64', min=0, max=1),
+        ),
+    ],
+)
+class TestBinomial(unittest.TestCase):
+    def setUp(self):
+        startup_program = paddle.static.Program()
+        main_program = paddle.static.Program()
+        executor = paddle.static.Executor(self.place)
+        with paddle.static.program_guard(main_program, startup_program):
+            probs = paddle.static.data(
+                'probs', self.probs.shape, self.probs.dtype
+            )
+            total_count = paddle.static.data(
+                'total_count', self.total_count.shape, self.total_count.dtype
+            )
+            dist = Binomial(total_count, probs)
+            mean = dist.mean
+            var = dist.variance
+            entropy = dist.entropy()
+            large_samples = dist.sample(shape=(1000,))
+        fetch_list = [mean, var, entropy, large_samples]
+        feed = {
+            'probs': self.probs,
+            'total_count': self.total_count,
+        }
+
+        executor.run(startup_program)
+        [
+            self.mean,
+            self.var,
+            self.entropy,
+            self.large_samples,
+        ] = executor.run(main_program, feed=feed, fetch_list=fetch_list)
+
+    def test_mean(self):
+        self.assertEqual(str(self.mean.dtype).split('.')[-1], self.probs.dtype)
+        np.testing.assert_allclose(
+            self.mean,
+            self._np_mean(),
+            rtol=config.RTOL.get(str(self.probs.dtype)),
+            atol=config.ATOL.get(str(self.probs.dtype)),
+        )
+
+    def test_variance(self):
+        self.assertEqual(str(self.var.dtype).split('.')[-1], self.probs.dtype)
+        np.testing.assert_allclose(
+            self.var,
+            self._np_variance(),
+            rtol=config.RTOL.get(str(self.probs.dtype)),
+            atol=config.ATOL.get(str(self.probs.dtype)),
+        )
+
+    def test_entropy(self):
+        self.assertEqual(
+            str(self.entropy.dtype).split('.')[-1], self.probs.dtype
+        )
+        np.testing.assert_allclose(
+            self.entropy,
+            self._np_entropy(),
+            rtol=config.RTOL.get(str(self.probs.dtype)),
+            atol=config.ATOL.get(str(self.probs.dtype)),
+        )
+
+    def test_sample(self):
+        self.assertEqual(
+            str(self.large_samples.dtype).split('.')[-1], self.probs.dtype
+        )
+        sample_mean = self.large_samples.mean(axis=0)
+        sample_variance = self.large_samples.var(axis=0)
+        np.testing.assert_allclose(sample_mean, self.mean, atol=0, rtol=0.20)
+        np.testing.assert_allclose(sample_variance, self.var, atol=0, rtol=0.20)
+
+    def _np_variance(self):
+        return scipy.stats.binom.var(self.total_count, self.probs)
+
+    def _np_mean(self):
+        return scipy.stats.binom.mean(self.total_count, self.probs)
+
+    def _np_entropy(self):
+        return scipy.stats.binom.entropy(self.total_count, self.probs)
+
+
+@parameterize.place(config.DEVICES)
+@parameterize.parameterize_cls(
+    (parameterize.TEST_CASE_NAME, 'total_count', 'probs', 'value'),
+    [
+        (
+            'zero-dim',
+            np.array(10),
+            np.array(0.6).astype('float64'),
+            np.array([2.0, 3.0, 5.0]).astype('float64'),
+        ),
+        (
+            'value-same-shape',
+            np.array([10]).astype('int64'),
+            np.array([0.2, 0.3, 0.5]).astype('float64'),
+            np.array([2.0, 3.0, 5.0]).astype('float64'),
+        ),
+        (
+            'value-broadcast-shape',
+            np.array([10]),
+            np.array([[0.3, 0.7], [0.5, 0.5]]),
+            np.array([[[4.0, 6.0], [8.0, 2.0]], [[2.0, 4.0], [9.0, 7.0]]]),
+        ),
+    ],
+)
+class TestBinomialProbs(unittest.TestCase):
+    def setUp(self):
+        startup_program = paddle.static.Program()
+        main_program = paddle.static.Program()
+        executor = paddle.static.Executor(self.place)
+
+        with paddle.static.program_guard(main_program, startup_program):
+            total_count = paddle.static.data(
+                'total_count', self.total_count.shape, self.total_count.dtype
+            )
+            probs = paddle.static.data(
+                'probs', self.probs.shape, self.probs.dtype
+            )
+            value = paddle.static.data(
+                'value', self.value.shape, self.value.dtype
+            )
+            dist = Binomial(total_count, probs)
+            pmf = dist.prob(value)
+        feed = {
+            'total_count': self.total_count,
+            'probs': self.probs,
+            'value': self.value,
+        }
+        fetch_list = [pmf]
+
+        executor.run(startup_program)
+        [self.pmf] = executor.run(
+            main_program, feed=feed, fetch_list=fetch_list
+        )
+
+    def test_prob(self):
+        np.testing.assert_allclose(
+            self.pmf,
+            scipy.stats.binom.pmf(self.value, self.total_count, self.probs),
+            rtol=config.RTOL.get(str(self.probs.dtype)),
+            atol=config.ATOL.get(str(self.probs.dtype)),
+        )
+
+
+@parameterize.place(config.DEVICES)
+@parameterize.parameterize_cls(
+    (parameterize.TEST_CASE_NAME, 'n_1', 'p_1', 'n_2', 'p_2'),
+    [
+        (
+            'multi-dim-probability',
+            np.array([32]),
+            parameterize.xrand((1, 2), dtype='float64', min=0, max=1),
+            np.array([32]),
+            parameterize.xrand((1, 2), dtype='float64', min=0, max=1),
+        ),
+    ],
+)
+class TestBinomialKL(unittest.TestCase):
+    def setUp(self):
+        startup_program = paddle.static.Program()
+        main_program = paddle.static.Program()
+        executor = paddle.static.Executor(self.place)
+
+        with paddle.static.program_guard(main_program, startup_program):
+            n_1 = paddle.static.data('n_1', self.n_1.shape, self.n_1.dtype)
+            p_1 = paddle.static.data('p_1', self.p_1.shape, self.p_1.dtype)
+            n_2 = paddle.static.data('n_2', self.n_2.shape, self.n_2.dtype)
+            p_2 = paddle.static.data('p_2', self.p_2.shape, self.p_2.dtype)
+            dist1 = Binomial(n_1, p_1)
+            dist2 = Binomial(n_2, p_2)
+            kl_dist1_dist2 = dist1.kl_divergence(dist2)
+        feed = {
+            'n_1': self.n_1,
+            'p_1': self.p_1,
+            'n_2': self.n_2,
+            'p_2': self.p_2,
+        }
+        fetch_list = [kl_dist1_dist2]
+
+        executor.run(startup_program)
+        [self.kl_dist1_dist2] = executor.run(
+            main_program, feed=feed, fetch_list=fetch_list
+        )
+
+    def test_kl_divergence(self):
+        kl0 = self.kl_dist1_dist2
+        kl1 = self.kl_divergence_scipy()
+
+        self.assertEqual(tuple(kl0.shape), self.p_1.shape)
+        self.assertEqual(tuple(kl1.shape), self.p_1.shape)
+        np.testing.assert_allclose(
+            kl0,
+            kl1,
+            rtol=config.RTOL.get(str(self.p_1.dtype)),
+            atol=config.ATOL.get(str(self.p_1.dtype)),
+        )
+
+    def kl_divergence_scipy(self):
+        support = np.arange(1 + self.n_1.max(), dtype=self.p_1.dtype)
+        support = support.reshape((-1,) + (1,) * len(self.p_1.shape))
+        log_prob_1 = scipy.stats.binom.logpmf(support, self.n_1, self.p_1)
+        log_prob_2 = scipy.stats.binom.logpmf(support, self.n_2, self.p_2)
+        return (np.exp(log_prob_1) * (log_prob_1 - log_prob_2)).sum(0)
+
+
+if __name__ == '__main__':
+    unittest.main()

test/distribution/test_distribution_categorical.py

@@ -0,0 +1,468 @@
+# Copyright (c) 2021 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
+#
+#     http://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 test_distribution import DistributionNumpy
+
+import paddle
+from paddle import base
+from paddle.distribution import Categorical, Distribution, Normal, Uniform
+
+np.random.seed(2022)
+
+
+class CategoricalNumpy(DistributionNumpy):
+    def __init__(self, logits):
+        self.logits = np.array(logits).astype('float32')
+
+    def entropy(self):
+        logits = self.logits - np.max(self.logits, axis=-1, keepdims=True)
+        e_logits = np.exp(logits)
+        z = np.sum(e_logits, axis=-1, keepdims=True)
+        prob = e_logits / z
+        return -1.0 * np.sum(prob * (logits - np.log(z)), axis=-1)
+
+    def kl_divergence(self, other):
+        logits = self.logits - np.max(self.logits, axis=-1, keepdims=True)
+        other_logits = other.logits - np.max(
+            other.logits, axis=-1, keepdims=True
+        )
+        e_logits = np.exp(logits)
+        other_e_logits = np.exp(other_logits)
+        z = np.sum(e_logits, axis=-1, keepdims=True)
+        other_z = np.sum(other_e_logits, axis=-1, keepdims=True)
+        prob = e_logits / z
+        return np.sum(
+            prob * (logits - np.log(z) - other_logits + np.log(other_z)),
+            axis=-1,
+            keepdims=True,
+        )
+
+
+class CategoricalTest(unittest.TestCase):
+    def setUp(self, use_gpu=False, batch_size=3, dims=5):
+        self.use_gpu = use_gpu
+        if not use_gpu:
+            self.place = base.CPUPlace()
+            self.gpu_id = -1
+        else:
+            self.place = base.CUDAPlace(0)
+            self.gpu_id = 0
+
+        self.batch_size = batch_size
+        self.dims = dims
+        self.init_numpy_data(batch_size, dims)
+
+        paddle.disable_static(self.place)
+        self.init_dynamic_data(batch_size, dims)
+
+        paddle.enable_static()
+        self.test_program = base.Program()
+        self.executor = base.Executor(self.place)
+        self.init_static_data(batch_size, dims)
+
+    def init_numpy_data(self, batch_size, dims):
+        # input logtis is 2-D Tensor
+        # value used in probs and log_prob method is 1-D Tensor
+        self.logits_np = np.random.rand(batch_size, dims).astype('float32')
+        self.other_logits_np = np.random.rand(batch_size, dims).astype(
+            'float32'
+        )
+        self.value_np = np.array([2, 1, 3]).astype('int64')
+
+        self.logits_shape = [batch_size, dims]
+        # dist_shape = logits_shape[:-1], it represents the number of
+        #  different distributions.
+        self.dist_shape = [batch_size]
+        # sample shape represents the number of samples
+        self.sample_shape = [2, 4]
+        # value used in probs and log_prob method
+        # If value is 1-D and logits is 2-D or higher dimension, value will be
+        #  broadcasted to have the same number of distributions with logits.
+        # If value is 2-D or higher dimentsion, it should have the same number
+        #  of distributions with logtis. ``value[:-1] = logits[:-1]
+        self.value_shape = [3]
+
+    def init_dynamic_data(self, batch_size, dims):
+        self.logits = paddle.to_tensor(self.logits_np)
+        self.other_logits = paddle.to_tensor(self.other_logits_np)
+        self.value = paddle.to_tensor(self.value_np)
+
+    def init_static_data(self, batch_size, dims):
+        with base.program_guard(self.test_program):
+            self.logits_static = paddle.static.data(
+                name='logits', shape=self.logits_shape, dtype='float32'
+            )
+            self.other_logits_static = paddle.static.data(
+                name='other_logits', shape=self.logits_shape, dtype='float32'
+            )
+            self.value_static = paddle.static.data(
+                name='value', shape=self.value_shape, dtype='int64'
+            )
+
+    def get_numpy_selected_probs(self, probability):
+        np_probs = np.zeros(self.dist_shape + self.value_shape)
+        for i in range(self.batch_size):
+            for j in range(3):
+                np_probs[i][j] = probability[i][self.value_np[j]]
+        return np_probs
+
+    def compare_with_numpy(self, fetch_list, tolerance=1e-6):
+        sample, entropy, kl, probs, log_prob = fetch_list
+        log_tolerance = 1e-4
+
+        np.testing.assert_equal(
+            sample.shape, self.sample_shape + self.dist_shape
+        )
+
+        np_categorical = CategoricalNumpy(self.logits_np)
+        np_other_categorical = CategoricalNumpy(self.other_logits_np)
+        np_entropy = np_categorical.entropy()
+        np_kl = np_categorical.kl_divergence(np_other_categorical)
+
+        np.testing.assert_allclose(
+            entropy, np_entropy, rtol=log_tolerance, atol=log_tolerance
+        )
+        np.testing.assert_allclose(
+            kl, np_kl, rtol=log_tolerance, atol=log_tolerance
+        )
+
+        sum_dist = np.sum(self.logits_np, axis=-1, keepdims=True)
+        probability = self.logits_np / sum_dist
+        np_probs = self.get_numpy_selected_probs(probability)
+        np_log_prob = np.log(np_probs)
+
+        np.testing.assert_allclose(
+            probs, np_probs, rtol=tolerance, atol=tolerance
+        )
+        np.testing.assert_allclose(
+            log_prob, np_log_prob, rtol=tolerance, atol=tolerance
+        )
+
+    def test_categorical_distribution_dygraph(self, tolerance=1e-6):
+        paddle.disable_static(self.place)
+        categorical = Categorical(self.logits)
+        other_categorical = Categorical(self.other_logits)
+
+        sample = categorical.sample(self.sample_shape).numpy()
+        entropy = categorical.entropy().numpy()
+        kl = categorical.kl_divergence(other_categorical).numpy()
+        probs = categorical.probs(self.value).numpy()
+        log_prob = categorical.log_prob(self.value).numpy()
+
+        fetch_list = [sample, entropy, kl, probs, log_prob]
+        self.compare_with_numpy(fetch_list)
+
+    def test_categorical_distribution_static(self, tolerance=1e-6):
+        paddle.enable_static()
+        with base.program_guard(self.test_program):
+            categorical = Categorical(self.logits_static)
+            other_categorical = Categorical(self.other_logits_static)
+
+            sample = categorical.sample(self.sample_shape)
+            entropy = categorical.entropy()
+            kl = categorical.kl_divergence(other_categorical)
+            probs = categorical.probs(self.value_static)
+            log_prob = categorical.log_prob(self.value_static)
+
+            fetch_list = [sample, entropy, kl, probs, log_prob]
+
+        feed_vars = {
+            'logits': self.logits_np,
+            'other_logits': self.other_logits_np,
+            'value': self.value_np,
+        }
+
+        self.executor.run(base.default_startup_program())
+        fetch_list = self.executor.run(
+            program=self.test_program, feed=feed_vars, fetch_list=fetch_list
+        )
+
+        self.compare_with_numpy(fetch_list)
+
+
+class CategoricalTest2(CategoricalTest):
+    def init_numpy_data(self, batch_size, dims):
+        # input logtis is 2-D Tensor with dtype Float64
+        # value used in probs and log_prob method is 1-D Tensor
+        self.logits_np = np.random.rand(batch_size, dims).astype('float64')
+        self.other_logits_np = np.random.rand(batch_size, dims).astype(
+            'float64'
+        )
+        self.value_np = np.array([2, 1, 3]).astype('int64')
+
+        self.logits_shape = [batch_size, dims]
+        self.dist_shape = [batch_size]
+        self.sample_shape = [2, 4]
+        self.value_shape = [3]
+
+    def init_static_data(self, batch_size, dims):
+        with base.program_guard(self.test_program):
+            self.logits_static = paddle.static.data(
+                name='logits', shape=self.logits_shape, dtype='float64'
+            )
+            self.other_logits_static = paddle.static.data(
+                name='other_logits', shape=self.logits_shape, dtype='float64'
+            )
+            self.value_static = paddle.static.data(
+                name='value', shape=self.value_shape, dtype='int64'
+            )
+
+
+class CategoricalTest3(CategoricalTest):
+    def init_dynamic_data(self, batch_size, dims):
+        # input logtis is 2-D numpy.ndarray with dtype Float32
+        # value used in probs and log_prob method is 1-D Tensor
+        self.logits = self.logits_np
+        self.other_logits = self.other_logits_np
+        self.value = paddle.to_tensor(self.value_np)
+
+    def init_static_data(self, batch_size, dims):
+        with base.program_guard(self.test_program):
+            self.logits_static = self.logits_np
+            self.other_logits_static = self.other_logits_np
+            self.value_static = paddle.static.data(
+                name='value', shape=self.value_shape, dtype='int64'
+            )
+
+
+class CategoricalTest4(CategoricalTest):
+    def init_numpy_data(self, batch_size, dims):
+        # input logtis is 2-D numpy.ndarray with dtype Float64
+        # value used in probs and log_prob method is 1-D Tensor
+        self.logits_np = np.random.rand(batch_size, dims).astype('float64')
+        self.other_logits_np = np.random.rand(batch_size, dims).astype(
+            'float64'
+        )
+        self.value_np = np.array([2, 1, 3]).astype('int64')
+
+        self.logits_shape = [batch_size, dims]
+        self.dist_shape = [batch_size]
+        self.sample_shape = [2, 4]
+        self.value_shape = [3]
+
+    def init_dynamic_data(self, batch_size, dims):
+        self.logits = self.logits_np
+        self.other_logits = self.other_logits_np
+        self.value = paddle.to_tensor(self.value_np)
+
+    def init_static_data(self, batch_size, dims):
+        with base.program_guard(self.test_program):
+            self.logits_static = self.logits_np
+            self.other_logits_static = self.other_logits_np
+            self.value_static = paddle.static.data(
+                name='value', shape=self.value_shape, dtype='int64'
+            )
+
+
+# test shape of logits and value used in probs and log_prob method
+class CategoricalTest5(CategoricalTest):
+    def init_numpy_data(self, batch_size, dims):
+        # input logtis is 1-D Tensor
+        # value used in probs and log_prob method is 1-D Tensor
+        self.logits_np = np.random.rand(dims).astype('float32')
+        self.other_logits_np = np.random.rand(dims).astype('float32')
+        self.value_np = np.array([2, 1, 3]).astype('int64')
+
+        self.logits_shape = [dims]
+        self.dist_shape = []
+        self.sample_shape = [2, 4]
+        self.value_shape = [3]
+
+    def get_numpy_selected_probs(self, probability):
+        np_probs = np.zeros(self.value_shape)
+        for i in range(3):
+            np_probs[i] = probability[self.value_np[i]]
+        return np_probs
+
+
+class CategoricalTest6(CategoricalTest):
+    def init_numpy_data(self, batch_size, dims):
+        # input logtis is 2-D Tensor
+        # value used in probs and log_prob method has the same number of batches with input
+        self.logits_np = np.random.rand(3, 5).astype('float32')
+        self.other_logits_np = np.random.rand(3, 5).astype('float32')
+        self.value_np = np.array([[2, 1], [0, 3], [2, 3]]).astype('int64')
+
+        self.logits_shape = [3, 5]
+        self.dist_shape = [3]
+        self.sample_shape = [2, 4]
+        self.value_shape = [3, 2]
+
+    def get_numpy_selected_probs(self, probability):
+        np_probs = np.zeros(self.value_shape)
+        for i in range(3):
+            for j in range(2):
+                np_probs[i][j] = probability[i][self.value_np[i][j]]
+        return np_probs
+
+
+class CategoricalTest7(CategoricalTest):
+    def init_numpy_data(self, batch_size, dims):
+        # input logtis is 3-D Tensor
+        # value used in probs and log_prob method has the same number of distributions with input
+        self.logits_np = np.random.rand(3, 2, 5).astype('float32')
+        self.other_logits_np = np.random.rand(3, 2, 5).astype('float32')
+        self.value_np = np.array([2, 1, 3]).astype('int64')
+
+        self.logits_shape = [3, 2, 5]
+        self.dist_shape = [3, 2]
+        self.sample_shape = [2, 4]
+        self.value_shape = [3]
+
+    def get_numpy_selected_probs(self, probability):
+        np_probs = np.zeros(self.dist_shape + self.value_shape)
+        for i in range(3):
+            for j in range(2):
+                for k in range(3):
+                    np_probs[i][j][k] = probability[i][j][self.value_np[k]]
+        return np_probs
+
+
+class CategoricalTest8(CategoricalTest):
+    def init_dynamic_data(self, batch_size, dims):
+        # input logtis is 2-D list
+        # value used in probs and log_prob method is 1-D Tensor
+        self.logits = self.logits_np.tolist()
+        self.other_logits = self.other_logits_np.tolist()
+        self.value = paddle.to_tensor(self.value_np)
+
+    def init_static_data(self, batch_size, dims):
+        with base.program_guard(self.test_program):
+            self.logits_static = self.logits_np.tolist()
+            self.other_logits_static = self.other_logits_np.tolist()
+            self.value_static = paddle.static.data(
+                name='value', shape=self.value_shape, dtype='int64'
+            )
+
+
+class CategoricalTest9(CategoricalTest):
+    def init_dynamic_data(self, batch_size, dims):
+        # input logtis is 2-D tuple
+        # value used in probs and log_prob method is 1-D Tensor
+        self.logits = tuple(self.logits_np.tolist())
+        self.other_logits = tuple(self.other_logits_np.tolist())
+        self.value = paddle.to_tensor(self.value_np)
+
+    def init_static_data(self, batch_size, dims):
+        with base.program_guard(self.test_program):
+            self.logits_static = tuple(self.logits_np.tolist())
+            self.other_logits_static = tuple(self.other_logits_np.tolist())
+            self.value_static = paddle.static.data(
+                name='value', shape=self.value_shape, dtype='int64'
+            )
+
+
+class DistributionTestError(unittest.TestCase):
+    def test_distribution_error(self):
+        distribution = Distribution()
+
+        self.assertRaises(NotImplementedError, distribution.sample)
+        self.assertRaises(NotImplementedError, distribution.entropy)
+
+        normal = Normal(0.0, 1.0)
+        self.assertRaises(
+            NotImplementedError, distribution.kl_divergence, normal
+        )
+
+        value_npdata = np.array([0.8], dtype="float32")
+        value_tensor = paddle.tensor.create_tensor(dtype="float32")
+        self.assertRaises(
+            NotImplementedError, distribution.log_prob, value_tensor
+        )
+        self.assertRaises(NotImplementedError, distribution.probs, value_tensor)
+
+    def test_normal_error(self):
+        paddle.enable_static()
+        normal = Normal(0.0, 1.0)
+
+        value = [1.0, 2.0]
+        # type of value must be variable
+        self.assertRaises(TypeError, normal.log_prob, value)
+
+        value = [1.0, 2.0]
+        # type of value must be variable
+        self.assertRaises(TypeError, normal.probs, value)
+
+        shape = 1.0
+        # type of shape must be list
+        self.assertRaises(TypeError, normal.sample, shape)
+
+        seed = 1.0
+        # type of seed must be int
+        self.assertRaises(TypeError, normal.sample, [2, 3], seed)
+
+        normal_other = Uniform(1.0, 2.0)
+        # type of other must be an instance of Normal
+        self.assertRaises(TypeError, normal.kl_divergence, normal_other)
+
+    def test_uniform_error(self):
+        paddle.enable_static()
+        uniform = Uniform(0.0, 1.0)
+
+        value = [1.0, 2.0]
+        # type of value must be variable
+        self.assertRaises(TypeError, uniform.log_prob, value)
+
+        value = [1.0, 2.0]
+        # type of value must be variable
+        self.assertRaises(TypeError, uniform.probs, value)
+
+        shape = 1.0
+        # type of shape must be list
+        self.assertRaises(TypeError, uniform.sample, shape)
+
+        seed = 1.0
+        # type of seed must be int
+        self.assertRaises(TypeError, uniform.sample, [2, 3], seed)
+
+    def test_categorical_error(self):
+        paddle.enable_static()
+
+        categorical = Categorical([0.4, 0.6])
+
+        value = [1, 0]
+        # type of value must be variable
+        self.assertRaises(AttributeError, categorical.log_prob, value)
+
+        value = [1, 0]
+        # type of value must be variable
+        self.assertRaises(AttributeError, categorical.probs, value)
+
+        shape = 1.0
+        # type of shape must be list
+        self.assertRaises(TypeError, categorical.sample, shape)
+
+        categorical_other = Uniform(1.0, 2.0)
+        # type of other must be an instance of Categorical
+        self.assertRaises(
+            TypeError, categorical.kl_divergence, categorical_other
+        )
+
+        def test_shape_not_match_error():
+            # shape of value must match shape of logits
+            # value_shape[:-1] == logits_shape[:-1]
+            paddle.disable_static()
+            logits = paddle.rand([3, 5])
+            cat = Categorical(logits)
+            value = paddle.to_tensor([[2, 1, 3], [3, 2, 1]], dtype='int64')
+            cat.log_prob(value)
+
+        self.assertRaises(ValueError, test_shape_not_match_error)
+
+
+if __name__ == '__main__':
+    unittest.main()

test/distribution/test_distribution_cauchy.py

@@ -0,0 +1,709 @@
+# Copyright (c) 2021 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
+#
+#     http://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
+import scipy.special
+import scipy.stats
+from distribution.config import ATOL, DEVICES, RTOL
+from parameterize import (
+    TEST_CASE_NAME,
+    parameterize_cls,
+    parameterize_func,
+    place,
+)
+from test_distribution import DistributionNumpy
+
+import paddle
+from paddle.base.data_feeder import convert_dtype
+from paddle.distribution import Cauchy
+from paddle.distribution.kl import kl_divergence
+
+np.random.seed(2023)
+paddle.seed(2023)
+
+
+def _kstest(samples_a, samples_b):
+    """Uses the Kolmogorov-Smirnov test for goodness of fit."""
+    _, p_value = scipy.stats.ks_2samp(samples_a, samples_b)
+    return not p_value < 0.005
+
+
+class CauchyNumpy(DistributionNumpy):
+    def __init__(self, loc, scale):
+        loc = np.array(loc)
+        scale = np.array(scale)
+        if str(loc.dtype) not in ['float32', 'float64']:
+            self.dtype = 'float32'
+        else:
+            self.dtype = loc.dtype
+
+        self.batch_shape = (loc + scale).shape
+
+        self.loc = loc.astype(self.dtype)
+        self.scale = scale.astype(self.dtype)
+        self.rv = scipy.stats.cauchy(loc=loc, scale=scale)
+
+    def sample(self, shape):
+        shape = np.array(shape, dtype='int')
+        if shape.ndim:
+            shape = shape.tolist()
+        else:
+            shape = [shape.tolist()]
+        return self.rv.rvs(size=shape + list(self.batch_shape))
+
+    def log_prob(self, value):
+        return self.rv.logpdf(value)
+
+    def prob(self, value):
+        return self.rv.pdf(value)
+
+    def cdf(self, value):
+        return self.rv.cdf(value)
+
+    def entropy(self):
+        return self.rv.entropy()
+
+    def kl_divergence(self, other):
+        a_loc = self.loc
+        b_loc = other.loc
+
+        a_scale = self.scale
+        b_scale = other.scale
+
+        t1 = np.log(np.power(a_scale + b_scale, 2) + np.power(a_loc - b_loc, 2))
+        t2 = np.log(4 * a_scale * b_scale)
+        return t1 - t2
+
+
+class CauchyTest(unittest.TestCase):
+    def setUp(self):
+        paddle.disable_static(self.place)
+        with paddle.base.dygraph.guard(self.place):
+            # just for convenience
+            self.dtype = self.expected_dtype
+
+            # init numpy with `dtype`
+            self.init_numpy_data(self.loc, self.scale, self.dtype)
+
+            # init paddle and check dtype convert.
+            self.init_dynamic_data(
+                self.loc, self.scale, self.default_dtype, self.dtype
+            )
+
+    def init_numpy_data(self, loc, scale, dtype):
+        loc = np.array(loc).astype(dtype)
+        scale = np.array(scale).astype(dtype)
+        self.rv_np = CauchyNumpy(loc=loc, scale=scale)
+
+    def init_dynamic_data(self, loc, scale, default_dtype, dtype):
+        self.rv_paddle = Cauchy(loc=loc, scale=scale)
+        self.assertTrue(
+            dtype == convert_dtype(self.rv_paddle.loc.dtype),
+            (dtype, self.rv_paddle.loc.dtype),
+        )
+        self.assertTrue(
+            dtype == convert_dtype(self.rv_paddle.scale.dtype),
+            (dtype, self.rv_paddle.scale.dtype),
+        )
+
+
+@place(DEVICES)
+@parameterize_cls(
+    (TEST_CASE_NAME, 'loc', 'scale', 'default_dtype', 'expected_dtype'),
+    [
+        # 0-D params
+        (
+            'params_0d_32_1',
+            paddle.full((), 0.1),
+            paddle.full((), 1.2),
+            'float32',
+            'float32',
+        ),
+        (
+            'params_0d_32_2',
+            paddle.full((), -1.2),
+            paddle.full((), 2.3),
+            'float32',
+            'float32',
+        ),
+        (
+            'params_0d_64_1',
+            paddle.full((), 0.1, dtype='float64'),
+            paddle.full((), 1.2, dtype='float64'),
+            'float64',
+            'float64',
+        ),
+        (
+            'params_0d_64_2',
+            paddle.full((), -1.2, dtype='float64'),
+            paddle.full((), 2.3, dtype='float64'),
+            'float64',
+            'float64',
+        ),
+        # 1-D params
+        ('params_float_1', 0.1, 1.2, 'float64', 'float32'),
+        ('params_float_2', -1.2, 2.3, 'float64', 'float32'),
+        (
+            'params_tensor_32_1',
+            paddle.to_tensor(0.1),
+            paddle.to_tensor(1.2),
+            'float32',
+            'float32',
+        ),
+        (
+            'params_tensor_32_2',
+            paddle.to_tensor(-1.2),
+            paddle.to_tensor(2.3),
+            'float32',
+            'float32',
+        ),
+        (
+            'params_tensor_64_1',
+            paddle.to_tensor(0.1, dtype='float64'),
+            paddle.to_tensor(1.2, dtype='float64'),
+            'float64',
+            'float64',
+        ),
+        (
+            'params_tensor_64_2',
+            paddle.to_tensor(-1.2, dtype='float64'),
+            paddle.to_tensor(2.3, dtype='float64'),
+            'float64',
+            'float64',
+        ),
+        (
+            'params_tensor_list',
+            paddle.to_tensor([0.1]),
+            paddle.to_tensor([1.2]),
+            'float32',
+            'float32',
+        ),
+        (
+            'params_tensor_tuple',
+            paddle.to_tensor((0.1,)),
+            paddle.to_tensor((1.2,)),
+            'float32',
+            'float32',
+        ),
+        # N-D params
+        (
+            'params_0d_1d_1',
+            paddle.full((), 0.1),
+            paddle.full((1,), 1.2),
+            'float32',
+            'float32',
+        ),
+        (
+            'params_0d_1d_2',
+            paddle.full((), 0.1),
+            paddle.to_tensor(1.2),
+            'float32',
+            'float32',
+        ),
+        (
+            'params_1d_0d_1',
+            paddle.full((1,), 0.1),
+            paddle.full((), 1.2),
+            'float32',
+            'float32',
+        ),
+        (
+            'params_1d_0d_2',
+            paddle.to_tensor(0.1),
+            paddle.full((), 1.2),
+            'float32',
+            'float32',
+        ),
+        (
+            'params_0d_3d',
+            paddle.full((), 0.1),
+            paddle.to_tensor([1.1, 2.2, 3.3]),
+            'float32',
+            'float32',
+        ),
+        (
+            'params_3d_0d',
+            paddle.to_tensor([0.1, -0.2, 0.3]),
+            paddle.full((), 1.2),
+            'float32',
+            'float32',
+        ),
+        (
+            'params_1d_3d',
+            paddle.full((1,), 0.1),
+            paddle.to_tensor([1.1, 2.2, 3.3]),
+            'float32',
+            'float32',
+        ),
+        (
+            'params_3d_1d',
+            paddle.to_tensor([0.1, -0.2, 0.3]),
+            paddle.full((1,), 1.2),
+            'float32',
+            'float32',
+        ),
+        (
+            'params_3d_3d',
+            paddle.to_tensor([0.1, -0.2, 0.3]),
+            paddle.to_tensor([1.1, 2.2, 3.3]),
+            'float32',
+            'float32',
+        ),
+    ],
+)
+class CauchyTestFeature(CauchyTest):
+    @parameterize_func(
+        [
+            (paddle.to_tensor([-0.3]),),
+            (paddle.to_tensor([0.3]),),
+            (paddle.to_tensor([1.3]),),
+            (paddle.to_tensor([5.3]),),
+            (paddle.to_tensor(0.3, dtype='float64'),),
+        ]
+    )
+    def test_log_prob(self, value):
+        with paddle.base.dygraph.guard(self.place):
+            if convert_dtype(value.dtype) == convert_dtype(
+                self.rv_paddle.loc.dtype
+            ):
+                log_prob = self.rv_paddle.log_prob(value)
+                np.testing.assert_allclose(
+                    log_prob,
+                    self.rv_np.log_prob(value),
+                    rtol=RTOL.get(self.dtype),
+                    atol=ATOL.get(self.dtype),
+                )
+                self.assertTrue(self.dtype == convert_dtype(log_prob.dtype))
+
+            else:
+                with self.assertWarns(UserWarning):
+                    self.rv_paddle.log_prob(value)
+
+    @parameterize_func(
+        [
+            (paddle.to_tensor([-0.3]),),
+            (paddle.to_tensor([0.3]),),
+            (paddle.to_tensor([1.3]),),
+            (paddle.to_tensor([5.3]),),
+            (paddle.to_tensor(0.3, dtype='float64'),),
+        ]
+    )
+    def test_prob(self, value):
+        with paddle.base.dygraph.guard(self.place):
+            if convert_dtype(value.dtype) == convert_dtype(
+                self.rv_paddle.loc.dtype
+            ):
+                prob = self.rv_paddle.prob(value)
+                np.testing.assert_allclose(
+                    prob,
+                    self.rv_np.prob(value),
+                    rtol=RTOL.get(self.dtype),
+                    atol=ATOL.get(self.dtype),
+                )
+                self.assertTrue(self.dtype == convert_dtype(prob.dtype))
+
+            else:
+                with self.assertWarns(UserWarning):
+                    self.rv_paddle.prob(value)
+
+    @parameterize_func(
+        [
+            (paddle.to_tensor([-0.3]),),
+            (paddle.to_tensor([0.3]),),
+            (paddle.to_tensor([1.3]),),
+            (paddle.to_tensor([5.3]),),
+            (paddle.to_tensor(0.3, dtype='float64'),),
+        ]
+    )
+    def test_cdf(self, value):
+        with paddle.base.dygraph.guard(self.place):
+            if convert_dtype(value.dtype) == convert_dtype(
+                self.rv_paddle.loc.dtype
+            ):
+                cdf = self.rv_paddle.cdf(value)
+                np.testing.assert_allclose(
+                    cdf,
+                    self.rv_np.cdf(value),
+                    rtol=RTOL.get(self.dtype),
+                    atol=ATOL.get(self.dtype),
+                )
+                self.assertTrue(self.dtype == convert_dtype(cdf.dtype))
+
+            else:
+                with self.assertWarns(UserWarning):
+                    self.rv_paddle.cdf(value)
+
+    def test_entropy(self):
+        with paddle.base.dygraph.guard(self.place):
+            np.testing.assert_allclose(
+                self.rv_paddle.entropy(),
+                self.rv_np.entropy(),
+                rtol=RTOL.get(self.dtype),
+                atol=ATOL.get(self.dtype),
+            )
+
+    @parameterize_func(
+        [
+            (0.6, 5.7),
+            (-0.6, 5.7),
+        ]
+    )
+    def test_kl_divergence(self, loc, scale):
+        with paddle.base.dygraph.guard(self.place):
+            # convert loc/scale to paddle's dtype(float32/float64)
+            rv_paddle_other = Cauchy(
+                loc=paddle.full((), loc, dtype=self.rv_paddle.loc.dtype),
+                scale=paddle.full((), scale, dtype=self.rv_paddle.scale.dtype),
+            )
+            rv_np_other = CauchyNumpy(loc=loc, scale=scale)
+
+            np.testing.assert_allclose(
+                self.rv_paddle.kl_divergence(rv_paddle_other),
+                self.rv_np.kl_divergence(rv_np_other),
+                rtol=RTOL.get(self.dtype),
+                atol=ATOL.get(self.dtype),
+            )
+
+            np.testing.assert_allclose(
+                kl_divergence(self.rv_paddle, rv_paddle_other),
+                self.rv_np.kl_divergence(rv_np_other),
+                rtol=RTOL.get(self.dtype),
+                atol=ATOL.get(self.dtype),
+            )
+
+
+@place(DEVICES)
+@parameterize_cls(
+    (
+        TEST_CASE_NAME,
+        'loc',
+        'scale',
+        'default_dtype',
+        'expected_dtype',
+        'shape',
+        'expected_shape',
+    ),
+    [
+        # 0-D params
+        (
+            'params_0d_0d_sample_1d',
+            paddle.full((), 0.1),
+            paddle.full((), 1.2),
+            'float32',
+            'float32',
+            [100],
+            [100],
+        ),
+        (
+            'params_0d_0d_sample_2d',
+            paddle.full((), 0.1),
+            paddle.full((), 1.2),
+            'float32',
+            'float32',
+            [100, 1],
+            [100, 1],
+        ),
+        (
+            'params_0d_0d_sample_3d',
+            paddle.full((), 0.1),
+            paddle.full((), 1.2),
+            'float32',
+            'float32',
+            [100, 2, 3],
+            [100, 2, 3],
+        ),
+        # 1-D params
+        (
+            'params_1d_1d_sample_1d_float',
+            0.1,
+            1.2,
+            'float64',
+            'float32',
+            paddle.to_tensor([100]),
+            [100],
+        ),
+        (
+            'params_1d_1d_sample_1d_32',
+            paddle.to_tensor([0.1]),
+            paddle.to_tensor([1.2]),
+            'float32',
+            'float32',
+            paddle.to_tensor([100]),
+            [100, 1],
+        ),
+        (
+            'params_1d_1d_sample_1d_64',
+            paddle.to_tensor([0.1], dtype='float64'),
+            paddle.to_tensor([1.2], dtype='float64'),
+            'float64',
+            'float64',
+            paddle.to_tensor([100]),
+            [100, 1],
+        ),
+        (
+            'params_1d_1d_sample_2d',
+            paddle.to_tensor([0.1]),
+            paddle.to_tensor([1.2]),
+            'float32',
+            'float32',
+            [100, 2],
+            [100, 2, 1],
+        ),
+        (
+            'params_1d_1d_sample_3d',
+            paddle.to_tensor([0.1]),
+            paddle.to_tensor([1.2]),
+            'float32',
+            'float32',
+            [100, 2, 3],
+            [100, 2, 3, 1],
+        ),
+        # N-D params
+        (
+            'params_0d_1d_sample_1d',
+            paddle.full((), 0.3),
+            paddle.to_tensor([1.2]),
+            'float32',
+            'float32',
+            [100],
+            [100, 1],
+        ),
+        (
+            'params_1d_0d_sample_1d',
+            paddle.to_tensor([0.3]),
+            paddle.full((), 1.2),
+            'float32',
+            'float32',
+            [100],
+            [100, 1],
+        ),
+        (
+            'params_0d_1d_sample_2d',
+            paddle.full((), 0.3),
+            paddle.to_tensor([1.2]),
+            'float32',
+            'float32',
+            [100, 2],
+            [100, 2, 1],
+        ),
+        (
+            'params_1d_0d_sample_2d',
+            paddle.to_tensor([0.3]),
+            paddle.full((), 1.2),
+            'float32',
+            'float32',
+            [100, 2],
+            [100, 2, 1],
+        ),
+        (
+            'params_1d_2d_sample_1d',
+            paddle.to_tensor([0.3]),
+            paddle.to_tensor((1.2, 2.3)),
+            'float32',
+            'float32',
+            [100],
+            [100, 2],
+        ),
+        (
+            'params_2d_1d_sample_1d',
+            paddle.to_tensor((0.3, -0.3)),
+            paddle.to_tensor([1.2]),
+            'float32',
+            'float32',
+            [100],
+            [100, 2],
+        ),
+        (
+            'params_2d_2d_sample_1d',
+            paddle.to_tensor((0.3, -0.3)),
+            paddle.to_tensor((1.2, 2.3)),
+            'float32',
+            'float32',
+            [100],
+            [100, 2],
+        ),
+        (
+            'params_2d_2d_sample_2d',
+            paddle.to_tensor((0.3, -0.3)),
+            paddle.to_tensor((1.2, 2.3)),
+            'float32',
+            'float32',
+            [100, 1],
+            [100, 1, 2],
+        ),
+        (
+            'params_1d_2d_sample_3d',
+            paddle.to_tensor([0.3]),
+            paddle.to_tensor((1.2, 2.3)),
+            'float32',
+            'float32',
+            [100, 1, 2],
+            [100, 1, 2, 2],
+        ),
+    ],
+)
+class CauchyTestSample(CauchyTest):
+    def test_sample(self):
+        with paddle.base.dygraph.guard(self.place):
+            sample_np = self.rv_np.sample(self.shape)
+            sample_paddle = self.rv_paddle.sample(self.shape)
+
+            self.assertEqual(list(sample_paddle.shape), self.expected_shape)
+            self.assertEqual(sample_paddle.dtype, self.rv_paddle.loc.dtype)
+
+            if len(self.expected_shape) > len(self.shape):
+                for i in range(self.expected_shape[-1]):
+                    self.assertTrue(
+                        _kstest(
+                            sample_np[..., i].reshape(-1),
+                            sample_paddle.numpy()[..., i].reshape(-1),
+                        )
+                    )
+            else:
+                self.assertTrue(
+                    _kstest(
+                        sample_np.reshape(-1),
+                        sample_paddle.numpy().reshape(-1),
+                    )
+                )
+
+    def test_rsample(self):
+        with paddle.base.dygraph.guard(self.place):
+            sample_np = self.rv_np.sample(self.shape)
+            rsample_paddle = self.rv_paddle.rsample(self.shape)
+
+            self.assertEqual(list(rsample_paddle.shape), self.expected_shape)
+            self.assertEqual(rsample_paddle.dtype, self.rv_paddle.loc.dtype)
+
+            if len(self.expected_shape) > len(self.shape):
+                for i in range(self.expected_shape[-1]):
+                    self.assertTrue(
+                        _kstest(
+                            sample_np[..., i].reshape(-1),
+                            rsample_paddle.numpy()[..., i].reshape(-1),
+                        )
+                    )
+            else:
+                self.assertTrue(
+                    _kstest(
+                        sample_np.reshape(-1),
+                        rsample_paddle.numpy().reshape(-1),
+                    )
+                )
+
+    def test_rsample_backpropagation(self):
+        with paddle.base.dygraph.guard(self.place):
+            self.rv_paddle.loc.stop_gradient = False
+            self.rv_paddle.scale.stop_gradient = False
+            rsample_paddle = self.rv_paddle.rsample(self.shape)
+            grads = paddle.grad(
+                [rsample_paddle], [self.rv_paddle.loc, self.rv_paddle.scale]
+            )
+            self.assertEqual(len(grads), 2)
+            self.assertEqual(grads[0].dtype, self.rv_paddle.loc.dtype)
+            self.assertEqual(grads[0].shape, self.rv_paddle.loc.shape)
+            self.assertEqual(grads[1].dtype, self.rv_paddle.scale.dtype)
+            self.assertEqual(grads[1].shape, self.rv_paddle.scale.shape)
+
+
+@place(DEVICES)
+@parameterize_cls([TEST_CASE_NAME], ['CauchyTestError'])
+class CauchyTestError(unittest.TestCase):
+    def setUp(self):
+        paddle.disable_static(self.place)
+
+    def test_bad_property(self):
+        """For property like mean/variance/stddev which is undefined in math,
+        we should raise `ValueError` instead of `NotImplementedError`.
+        """
+        with paddle.base.dygraph.guard(self.place):
+            rv = Cauchy(loc=0.0, scale=1.0)
+            with self.assertRaises(ValueError):
+                _ = rv.mean
+            with self.assertRaises(ValueError):
+                _ = rv.variance
+            with self.assertRaises(ValueError):
+                _ = rv.stddev
+
+    @parameterize_func(
+        [
+            (100,),  # int
+            (100.0,),  # float
+        ]
+    )
+    def test_bad_sample_shape_type(self, shape):
+        with paddle.base.dygraph.guard(self.place):
+            rv = Cauchy(loc=0.0, scale=1.0)
+
+            with self.assertRaises(TypeError):
+                _ = rv.sample(shape)
+
+            with self.assertRaises(TypeError):
+                _ = rv.rsample(shape)
+
+    @parameterize_func(
+        [
+            (1,),  # int
+            (1.0,),  # float
+            ([1.0],),  # list
+            ((1.0),),  # tuple
+            (np.array(1.0),),  # ndarray
+        ]
+    )
+    def test_bad_value_type(self, value):
+        with paddle.base.dygraph.guard(self.place):
+            rv = Cauchy(loc=0.0, scale=1.0)
+
+            with self.assertRaises(TypeError):
+                _ = rv.log_prob(value)
+
+            with self.assertRaises(TypeError):
+                _ = rv.prob(value)
+
+            with self.assertRaises(TypeError):
+                _ = rv.cdf(value)
+
+    @parameterize_func(
+        [
+            (np.array(1.0),),  # ndarray or other distribution
+        ]
+    )
+    def test_bad_kl_other_type(self, other):
+        with paddle.base.dygraph.guard(self.place):
+            rv = Cauchy(loc=0.0, scale=1.0)
+
+            with self.assertRaises(TypeError):
+                _ = rv.kl_divergence(other)
+
+    @parameterize_func(
+        [
+            (
+                paddle.to_tensor([0.1, 0.2]),
+                paddle.to_tensor([0.3, 0.4]),
+                paddle.to_tensor([0.1, 0.2, 0.3]),
+            ),
+        ]
+    )
+    def test_bad_broadcast(self, loc, scale, value):
+        with paddle.base.dygraph.guard(self.place):
+            rv = Cauchy(loc=loc, scale=scale)
+            self.assertRaises(ValueError, rv.cdf, value)
+            self.assertRaises(ValueError, rv.log_prob, value)
+            self.assertRaises(ValueError, rv.prob, value)
+
+
+if __name__ == '__main__':
+    unittest.main()

test/distribution/test_distribution_cauchy_static.py

@@ -0,0 +1,477 @@
+# Copyright (c) 2021 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
+#
+#     http://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 sys
+import unittest
+
+import numpy as np
+from distribution.config import ATOL, DEVICES, RTOL
+from parameterize import (
+    TEST_CASE_NAME,
+    parameterize_cls,
+    parameterize_func,
+    place,
+)
+
+sys.path.append("../../distribution")
+from test_distribution_cauchy import CauchyNumpy, _kstest
+
+import paddle
+from paddle.distribution import Cauchy
+from paddle.distribution.kl import kl_divergence
+
+np.random.seed(2023)
+paddle.seed(2023)
+paddle.enable_static()
+default_dtype = paddle.get_default_dtype()
+
+
+@place(DEVICES)
+@parameterize_cls(
+    (TEST_CASE_NAME, 'params'),
+    # params: name, loc, scale, loc_other, scale_other, value
+    [
+        (
+            'params',
+            (
+                # 1-D params
+                (
+                    'params_not_iterable',
+                    0.3,
+                    1.2,
+                    -1.2,
+                    2.3,
+                    3.4,
+                ),
+                (
+                    'params_not_iterable_and_broadcast_for_value',
+                    0.3,
+                    1.2,
+                    -1.2,
+                    2.3,
+                    np.array([[0.1, 1.2], [1.2, 3.4]], dtype=default_dtype),
+                ),
+                # N-D params
+                (
+                    'params_tuple_0305',
+                    (0.3, 0.5),
+                    0.7,
+                    -1.2,
+                    2.3,
+                    3.4,
+                ),
+                (
+                    'params_tuple_03050104',
+                    ((0.3, 0.5), (0.1, 0.4)),
+                    0.7,
+                    -1.2,
+                    2.3,
+                    3.4,
+                ),
+            ),
+        )
+    ],
+    test_pir=True,
+)
+class CauchyTestFeature(unittest.TestCase):
+    def setUp(self):
+        self.program = paddle.static.Program()
+        self.executor = paddle.static.Executor(self.place)
+
+        self.params_len = len(self.params)
+
+        with paddle.static.program_guard(self.program):
+            self.init_numpy_data(self.params)
+            self.init_static_data(self.params)
+
+    def init_numpy_data(self, params):
+        self.log_prob_np = []
+        self.prob_np = []
+        self.cdf_np = []
+        self.entropy_np = []
+        self.kl_np = []
+        self.shapes = []
+
+        for _, loc, scale, loc_other, scale_other, value in params:
+            rv_np = CauchyNumpy(loc=loc, scale=scale)
+            rv_np_other = CauchyNumpy(loc=loc_other, scale=scale_other)
+
+            self.log_prob_np.append(rv_np.log_prob(value))
+            self.prob_np.append(rv_np.prob(value))
+            self.cdf_np.append(rv_np.cdf(value))
+            self.entropy_np.append(rv_np.entropy())
+            self.kl_np.append(rv_np.kl_divergence(rv_np_other))
+            # paddle return data ndim>0
+            self.shapes.append(
+                (np.array(loc) + np.array(scale) + np.array(value)).shape
+                or (1,)
+            )
+
+    def init_static_data(self, params):
+        with paddle.static.program_guard(self.program):
+            rv_paddles = []
+            rv_paddles_other = []
+            values = []
+            for name, loc, scale, loc_other, scale_other, value in params:
+                if not isinstance(value, np.ndarray):
+                    value = paddle.full([1], value, dtype=default_dtype)
+                else:
+                    value = paddle.to_tensor(value, place=self.place)
+
+                # We should set name in static mode, or the executor confuse rv_paddles[i].
+                rv_paddles.append(
+                    Cauchy(
+                        loc=paddle.to_tensor(loc),
+                        scale=paddle.to_tensor(scale),
+                        name=name,
+                    )
+                )
+                rv_paddles_other.append(
+                    Cauchy(
+                        loc=paddle.to_tensor(loc_other),
+                        scale=paddle.to_tensor(scale_other),
+                        name=name,
+                    )
+                )
+                values.append(value)
+
+            results = self.executor.run(
+                self.program,
+                feed={},
+                fetch_list=[
+                    [
+                        rv_paddles[i].log_prob(values[i]),
+                        rv_paddles[i].prob(values[i]),
+                        rv_paddles[i].cdf(values[i]),
+                        rv_paddles[i].entropy(),
+                        rv_paddles[i].kl_divergence(rv_paddles_other[i]),
+                        kl_divergence(rv_paddles[i], rv_paddles_other[i]),
+                    ]
+                    for i in range(self.params_len)
+                ],
+            )
+
+            self.log_prob_paddle = []
+            self.prob_paddle = []
+            self.cdf_paddle = []
+            self.entropy_paddle = []
+            self.kl_paddle = []
+            self.kl_func_paddle = []
+            for i in range(self.params_len):
+                (
+                    _log_prob,
+                    _prob,
+                    _cdf,
+                    _entropy,
+                    _kl,
+                    _kl_func,
+                ) = results[i * 6 : (i + 1) * 6]
+                self.log_prob_paddle.append(_log_prob)
+                self.prob_paddle.append(_prob)
+                self.cdf_paddle.append(_cdf)
+                self.entropy_paddle.append(_entropy)
+                self.kl_paddle.append(_kl)
+                self.kl_func_paddle.append(_kl_func)
+
+    def test_all(self):
+        for i in range(self.params_len):
+            self._test_log_prob(i)
+            self._test_prob(i)
+            self._test_cdf(i)
+            self._test_entropy(i)
+            self._test_kl_divergence(i)
+
+    def _test_log_prob(self, i):
+        np.testing.assert_allclose(
+            self.log_prob_np[i],
+            self.log_prob_paddle[i],
+            rtol=RTOL.get(default_dtype),
+            atol=ATOL.get(default_dtype),
+        )
+
+        # check shape
+        self.assertTrue(self.log_prob_paddle[i].shape == self.shapes[i])
+
+    def _test_prob(self, i):
+        np.testing.assert_allclose(
+            self.prob_np[i],
+            self.prob_paddle[i],
+            rtol=RTOL.get(default_dtype),
+            atol=ATOL.get(default_dtype),
+        )
+
+        # check shape
+        self.assertTrue(self.prob_paddle[i].shape == self.shapes[i])
+
+    def _test_cdf(self, i):
+        np.testing.assert_allclose(
+            self.cdf_np[i],
+            self.cdf_paddle[i],
+            rtol=RTOL.get(default_dtype),
+            atol=ATOL.get(default_dtype),
+        )
+
+        # check shape
+        self.assertTrue(self.cdf_paddle[i].shape == self.shapes[i])
+
+    def _test_entropy(self, i):
+        np.testing.assert_allclose(
+            self.entropy_np[i],
+            self.entropy_paddle[i],
+            rtol=RTOL.get(default_dtype),
+            atol=ATOL.get(default_dtype),
+        )
+
+    def _test_kl_divergence(self, i):
+        np.testing.assert_allclose(
+            self.kl_np[i],
+            self.kl_paddle[i],
+            rtol=RTOL.get(default_dtype),
+            atol=ATOL.get(default_dtype),
+        )
+
+        np.testing.assert_allclose(
+            self.kl_np[i],
+            self.kl_func_paddle[i],
+            rtol=RTOL.get(default_dtype),
+            atol=ATOL.get(default_dtype),
+        )
+
+
+@place(DEVICES)
+@parameterize_cls(
+    (
+        TEST_CASE_NAME,
+        'loc',
+        'scale',
+        'shape',
+        'expected_shape',
+    ),
+    [
+        # 1-D params
+        (
+            'params_1d',
+            [0.1],
+            [1.2],
+            [100],
+            [100, 1],
+        ),
+        # N-D params
+        (
+            'params_2d',
+            [0.3],
+            [1.2, 2.3],
+            [100],
+            [100, 2],
+        ),
+    ],
+    test_pir=True,
+)
+class CauchyTestSample(unittest.TestCase):
+    def setUp(self):
+        self.program = paddle.static.Program()
+        self.executor = paddle.static.Executor(self.place)
+
+        with paddle.static.program_guard(self.program):
+            self.init_numpy_data(self.loc, self.scale, self.shape)
+            self.init_static_data(self.loc, self.scale, self.shape)
+
+    def init_numpy_data(self, loc, scale, shape):
+        self.rv_np = CauchyNumpy(loc=loc, scale=scale)
+        self.sample_np = self.rv_np.sample(shape)
+
+    def init_static_data(self, loc, scale, shape):
+        with paddle.static.program_guard(self.program):
+            self.rv_paddle = Cauchy(
+                loc=paddle.to_tensor(loc),
+                scale=paddle.to_tensor(scale),
+            )
+
+            [self.sample_paddle, self.rsample_paddle] = self.executor.run(
+                self.program,
+                feed={},
+                fetch_list=[
+                    self.rv_paddle.sample(shape),
+                    self.rv_paddle.rsample(shape),
+                ],
+            )
+
+    def test_sample(self):
+        with paddle.static.program_guard(self.program):
+            self.assertEqual(
+                list(self.sample_paddle.shape), self.expected_shape
+            )
+
+            for i in range(self.expected_shape[-1]):
+                self.assertTrue(
+                    _kstest(
+                        self.sample_np[..., i].reshape(-1),
+                        self.sample_paddle[..., i].reshape(-1),
+                    )
+                )
+
+    def test_rsample(self):
+        """Compare two samples from `rsample` method, one from scipy and another from paddle."""
+        with paddle.static.program_guard(self.program):
+            self.assertEqual(
+                list(self.rsample_paddle.shape), self.expected_shape
+            )
+
+            for i in range(self.expected_shape[-1]):
+                self.assertTrue(
+                    _kstest(
+                        self.sample_np[..., i].reshape(-1),
+                        self.rsample_paddle[..., i].reshape(-1),
+                    )
+                )
+
+
+@place(DEVICES)
+@parameterize_cls([TEST_CASE_NAME], ['CauchyTestError'], test_pir=True)
+class CauchyTestError(unittest.TestCase):
+    def setUp(self):
+        self.program = paddle.static.Program()
+        self.executor = paddle.static.Executor(self.place)
+
+    @parameterize_func(
+        [
+            ((0.3,),),  # tuple
+            ([0.3],),  # list
+            (np.array([0.3]),),  # ndarray
+            (-1j + 1,),  # complex
+            ('0',),  # str
+        ]
+    )
+    def test_bad_init_type(self, param):
+        """Test bad init for loc/scale"""
+        with paddle.static.program_guard(self.program):
+            with self.assertRaises(TypeError):
+                [_] = self.executor.run(
+                    self.program,
+                    feed={},
+                    fetch_list=[Cauchy(loc=0.0, scale=param).scale],
+                )
+            with self.assertRaises(TypeError):
+                [_] = self.executor.run(
+                    self.program,
+                    feed={},
+                    fetch_list=[Cauchy(loc=param, scale=1.0).loc],
+                )
+
+    def test_bad_property(self):
+        """For property like mean/variance/stddev which is undefined in math,
+        we should raise `ValueError` instead of `NotImplementedError`.
+        """
+        with paddle.static.program_guard(self.program):
+            rv = Cauchy(loc=0.0, scale=1.0)
+            with self.assertRaises(ValueError):
+                _ = rv.mean
+            with self.assertRaises(ValueError):
+                _ = rv.variance
+            with self.assertRaises(ValueError):
+                _ = rv.stddev
+
+    @parameterize_func(
+        [
+            (100,),  # int
+            (100.0,),  # float
+        ]
+    )
+    def test_bad_sample_shape_type(self, shape):
+        with paddle.static.program_guard(self.program):
+            rv = Cauchy(loc=0.0, scale=1.0)
+
+            with self.assertRaises(TypeError):
+                [_] = self.executor.run(
+                    self.program, feed={}, fetch_list=[rv.sample(shape)]
+                )
+
+            with self.assertRaises(TypeError):
+                [_] = self.executor.run(
+                    self.program, feed={}, fetch_list=[rv.rsample(shape)]
+                )
+
+    @parameterize_func(
+        [
+            (1,),  # int
+            (1.0,),  # float
+            ([1.0],),  # list
+            ((1.0),),  # tuple
+            (np.array(1.0),),  # ndarray
+        ]
+    )
+    def test_bad_value_type(self, value):
+        with paddle.static.program_guard(self.program):
+            rv = Cauchy(loc=0.0, scale=1.0)
+
+            with self.assertRaises(TypeError):
+                [_] = self.executor.run(
+                    self.program, feed={}, fetch_list=[rv.log_prob(value)]
+                )
+
+            with self.assertRaises(TypeError):
+                [_] = self.executor.run(
+                    self.program, feed={}, fetch_list=[rv.prob(value)]
+                )
+
+            with self.assertRaises(TypeError):
+                [_] = self.executor.run(
+                    self.program, feed={}, fetch_list=[rv.cdf(value)]
+                )
+
+    @parameterize_func(
+        [
+            (np.array(1.0),),  # ndarray or other distribution
+        ]
+    )
+    def test_bad_kl_other_type(self, other):
+        with paddle.static.program_guard(self.program):
+            rv = Cauchy(loc=0.0, scale=1.0)
+
+            with self.assertRaises(TypeError):
+                [_] = self.executor.run(
+                    self.program, feed={}, fetch_list=[rv.kl_divergence(other)]
+                )
+
+    @parameterize_func(
+        [
+            (paddle.to_tensor([0.1, 0.2, 0.3]),),
+        ]
+    )
+    def test_bad_broadcast(self, value):
+        with paddle.static.program_guard(self.program):
+            rv = Cauchy(
+                loc=paddle.to_tensor(0.0), scale=paddle.to_tensor((1.0, 2.0))
+            )
+
+            # `logits, value = paddle.broadcast_tensors([self.logits, value])`
+            # raise ValueError in dygraph, raise TypeError in static.
+            with self.assertRaises((TypeError, ValueError)):
+                [_] = self.executor.run(
+                    self.program, feed={}, fetch_list=[rv.cdf(value)]
+                )
+
+            with self.assertRaises((TypeError, ValueError)):
+                [_] = self.executor.run(
+                    self.program, feed={}, fetch_list=[rv.log_prob(value)]
+                )
+
+            with self.assertRaises((TypeError, ValueError)):
+                [_] = self.executor.run(
+                    self.program, feed={}, fetch_list=[rv.prob(value)]
+                )
+
+
+if __name__ == '__main__':
+    unittest.main()

test/distribution/test_distribution_chi2.py

@@ -0,0 +1,234 @@
+# Copyright (c) 2024 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
+#
+#     http://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 numbers
+import unittest
+
+import numpy as np
+import parameterize
+import scipy.stats
+from distribution import config
+
+import paddle
+from paddle.distribution import chi2
+
+np.random.seed(2024)
+paddle.seed(2024)
+
+
+@parameterize.place(config.DEVICES)
+@parameterize.parameterize_cls(
+    (parameterize.TEST_CASE_NAME, 'df'),
+    [
+        (
+            'one-dim',
+            parameterize.xrand(
+                (4,),
+                dtype='float32',
+                min=np.finfo(dtype='float32').tiny,
+            ),
+        ),
+        (
+            'multi-dim',
+            parameterize.xrand(
+                (2, 2),
+                dtype='float32',
+                min=np.finfo(dtype='float32').tiny,
+            ),
+        ),
+        (
+            'broadcast',
+            parameterize.xrand(
+                (2, 1),
+                dtype='float32',
+                min=np.finfo(dtype='float32').tiny,
+            ),
+        ),
+    ],
+)
+class TestChi2(unittest.TestCase):
+    def setUp(self):
+        df = self.df
+        if not isinstance(self.df, numbers.Real):
+            df = paddle.to_tensor(self.df)
+
+        self._paddle_chi2 = chi2.Chi2(df)
+
+    def test_mean(self):
+        with paddle.base.dygraph.guard(self.place):
+            np.testing.assert_allclose(
+                self._paddle_chi2.mean,
+                scipy.stats.chi2.mean(self.df),
+                rtol=config.RTOL.get(str(self._paddle_chi2.df.numpy().dtype)),
+                atol=config.ATOL.get(str(self._paddle_chi2.df.numpy().dtype)),
+            )
+
+    def test_variance(self):
+        with paddle.base.dygraph.guard(self.place):
+            np.testing.assert_allclose(
+                self._paddle_chi2.variance,
+                scipy.stats.chi2.var(self.df),
+                rtol=config.RTOL.get(str(self._paddle_chi2.df.numpy().dtype)),
+                atol=config.ATOL.get(str(self._paddle_chi2.df.numpy().dtype)),
+            )
+
+    def test_entropy(self):
+        with paddle.base.dygraph.guard(self.place):
+            np.testing.assert_allclose(
+                self._paddle_chi2.entropy(),
+                scipy.stats.chi2.entropy(self.df),
+                rtol=config.RTOL.get(str(self.df.dtype)),
+                atol=config.ATOL.get(str(self.df.dtype)),
+            )
+
+    def test_prob(self):
+        value = np.random.rand(*self._paddle_chi2.df.shape)
+        with paddle.base.dygraph.guard(self.place):
+            np.testing.assert_allclose(
+                self._paddle_chi2.prob(paddle.to_tensor(value)),
+                scipy.stats.chi2.pdf(value, self.df),
+                rtol=config.RTOL.get(str(self.df.dtype)),
+                atol=config.ATOL.get(str(self.df.dtype)),
+            )
+
+    def test_log_prob(self):
+        value = np.random.rand(*self._paddle_chi2.df.shape)
+        with paddle.base.dygraph.guard(self.place):
+            np.testing.assert_allclose(
+                self._paddle_chi2.log_prob(paddle.to_tensor(value)),
+                scipy.stats.chi2.logpdf(value, self.df),
+                rtol=config.RTOL.get(str(self.df.dtype)),
+                atol=config.ATOL.get(str(self.df.dtype)),
+            )
+
+
+@parameterize.place(config.DEVICES)
+@parameterize.parameterize_cls(
+    (parameterize.TEST_CASE_NAME, 'df'),
+    [
+        (
+            'one-dim',
+            parameterize.xrand(
+                (2,),
+                dtype='float32',
+                min=np.finfo(dtype='float32').tiny,
+            ),
+        ),
+        (
+            'multi-dim',
+            parameterize.xrand(
+                (2, 2),
+                dtype='float32',
+                min=np.finfo(dtype='float32').tiny,
+            ),
+        ),
+    ],
+)
+class TestChi2Sample(unittest.TestCase):
+    def setUp(self):
+        df = self.df
+        if not isinstance(self.df, numbers.Real):
+            df = paddle.to_tensor(self.df)
+
+        self._paddle_chi2 = chi2.Chi2(df)
+
+    def test_sample_shape(self):
+        cases = [
+            {
+                'input': (),
+                'expect': tuple(paddle.squeeze(self._paddle_chi2.df).shape),
+            },
+            {
+                'input': (2, 2),
+                'expect': (2, 2, *paddle.squeeze(self._paddle_chi2.df).shape),
+            },
+        ]
+        for case in cases:
+            self.assertTrue(
+                tuple(self._paddle_chi2.sample(case.get('input')).shape)
+                == case.get('expect')
+            )
+
+    def test_sample(self):
+        sample_shape = (30000,)
+        samples = self._paddle_chi2.sample(sample_shape)
+        sample_values = samples.numpy()
+
+        np.testing.assert_allclose(
+            sample_values.mean(axis=0),
+            scipy.stats.chi2.mean(self.df),
+            rtol=0.1,
+            atol=config.ATOL.get(str(self._paddle_chi2.df.numpy().dtype)),
+        )
+        np.testing.assert_allclose(
+            sample_values.var(axis=0),
+            scipy.stats.chi2.var(self.df),
+            rtol=0.1,
+            atol=config.ATOL.get(str(self._paddle_chi2.df.numpy().dtype)),
+        )
+
+
+@parameterize.place(config.DEVICES)
+@parameterize.parameterize_cls(
+    (parameterize.TEST_CASE_NAME, 'df'),
+    [
+        ('0-dim', 0.4),
+    ],
+)
+class TestChi2SampleKS(unittest.TestCase):
+    def setUp(self):
+        df = self.df
+        if not isinstance(self.df, numbers.Real):
+            df = paddle.to_tensor(self.df)
+
+        self._paddle_chi2 = chi2.Chi2(df)
+
+    def test_sample_ks(self):
+        sample_shape = (15000,)
+        samples = self._paddle_chi2.sample(sample_shape)
+        self.assertTrue(self._kstest(samples))
+
+    def _kstest(self, samples):
+        # Uses the Kolmogorov-Smirnov test for goodness of fit.
+        ks, _ = scipy.stats.kstest(samples, scipy.stats.chi2(self.df).cdf)
+        return ks < 0.02
+
+
+@parameterize.place(config.DEVICES)
+@parameterize.parameterize_cls(
+    (parameterize.TEST_CASE_NAME),
+    [
+        ('chi2_test_err'),
+    ],
+)
+class Chi2TestError(unittest.TestCase):
+    @parameterize.parameterize_func(
+        [
+            (-1.0, ValueError),  # df < 0
+            ((1.0, -1.0), ValueError),  # df < 0
+        ]
+    )
+    def test_bad_parameter(self, df, error):
+        with paddle.base.dygraph.guard(self.place):
+            self.assertRaises(error, chi2.Chi2, df)
+
+    @parameterize.parameterize_func([(10,)])  # not sequence object sample shape
+    def test_bad_sample_shape(self, shape):
+        with paddle.base.dygraph.guard(self.place):
+            _chi2 = chi2.Chi2(1.0)
+            self.assertRaises(TypeError, _chi2.sample, shape)
+
+
+if __name__ == '__main__':
+    unittest.main()

test/distribution/test_distribution_chi2_static.py

@@ -0,0 +1,285 @@
+# Copyright (c) 2024 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
+#
+#     http://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
+import parameterize
+import scipy.stats
+from distribution import config
+
+import paddle
+from paddle.distribution import chi2
+
+paddle.enable_static()
+
+np.random.seed(2024)
+paddle.seed(2024)
+
+
+@parameterize.place(config.DEVICES)
+@parameterize.parameterize_cls(
+    (parameterize.TEST_CASE_NAME, 'df'),
+    [
+        (
+            'one-dim',
+            parameterize.xrand(
+                (4,),
+                dtype='float32',
+                min=np.finfo(dtype='float32').tiny,
+            ),
+        ),
+        (
+            'multi-dim',
+            parameterize.xrand(
+                (2, 2),
+                dtype='float32',
+                min=np.finfo(dtype='float32').tiny,
+            ),
+        ),
+        (
+            'broadcast',
+            parameterize.xrand(
+                (2, 1),
+                dtype='float32',
+                min=np.finfo(dtype='float32').tiny,
+            ),
+        ),
+    ],
+)
+class TestChi2(unittest.TestCase):
+    def setUp(self):
+        self.program = paddle.static.Program()
+        self.executor = paddle.static.Executor(self.place)
+        with paddle.static.program_guard(self.program):
+            df = paddle.static.data('df', self.df.shape, self.df.dtype)
+            self._paddle_chi2 = chi2.Chi2(df)
+            self.feeds = {'df': self.df}
+
+    def test_mean(self):
+        with paddle.static.program_guard(self.program):
+            [mean] = self.executor.run(
+                self.program,
+                feed=self.feeds,
+                fetch_list=[self._paddle_chi2.mean],
+            )
+            np.testing.assert_allclose(
+                mean,
+                scipy.stats.chi2.mean(self.df),
+                rtol=config.RTOL.get(str(self.df.dtype)),
+                atol=config.ATOL.get(str(self.df.dtype)),
+            )
+
+    def test_variance(self):
+        with paddle.static.program_guard(self.program):
+            [variance] = self.executor.run(
+                self.program,
+                feed=self.feeds,
+                fetch_list=[self._paddle_chi2.variance],
+            )
+            np.testing.assert_allclose(
+                variance,
+                scipy.stats.chi2.var(self.df),
+                rtol=config.RTOL.get(str(self.df.dtype)),
+                atol=config.ATOL.get(str(self.df.dtype)),
+            )
+
+    def test_entropy(self):
+        with paddle.static.program_guard(self.program):
+            [entropy] = self.executor.run(
+                self.program,
+                feed=self.feeds,
+                fetch_list=[self._paddle_chi2.entropy()],
+            )
+            np.testing.assert_allclose(
+                entropy,
+                scipy.stats.chi2.entropy(self.df),
+                rtol=config.RTOL.get(str(self.df.dtype)),
+                atol=config.ATOL.get(str(self.df.dtype)),
+            )
+
+    def test_prob(self):
+        with paddle.static.program_guard(self.program):
+            value = paddle.static.data(
+                'value',
+                self._paddle_chi2.df.shape,
+                self._paddle_chi2.df.dtype,
+            )
+            prob = self._paddle_chi2.prob(value)
+
+            random_number = np.random.rand(*self._paddle_chi2.df.shape).astype(
+                self.df.dtype
+            )
+            feeds = dict(self.feeds, value=random_number)
+            [prob] = self.executor.run(
+                self.program, feed=feeds, fetch_list=[prob]
+            )
+            np.testing.assert_allclose(
+                prob,
+                scipy.stats.chi2.pdf(random_number, self.df),
+                rtol=config.RTOL.get(str(self.df.dtype)),
+                atol=config.ATOL.get(str(self.df.dtype)),
+            )
+
+    def test_log_prob(self):
+        with paddle.static.program_guard(self.program):
+            value = paddle.static.data(
+                'value',
+                self._paddle_chi2.df.shape,
+                self._paddle_chi2.df.dtype,
+            )
+            log_prob = self._paddle_chi2.log_prob(value)
+
+            random_number = np.random.rand(*self._paddle_chi2.df.shape).astype(
+                self.df.dtype
+            )
+            feeds = dict(self.feeds, value=random_number)
+            [log_prob] = self.executor.run(
+                self.program, feed=feeds, fetch_list=[log_prob]
+            )
+            np.testing.assert_allclose(
+                log_prob,
+                scipy.stats.chi2.logpdf(random_number, self.df),
+                rtol=config.RTOL.get(str(self.df.dtype)),
+                atol=config.ATOL.get(str(self.df.dtype)),
+            )
+
+
+@parameterize.place(config.DEVICES)
+@parameterize.parameterize_cls(
+    (parameterize.TEST_CASE_NAME, 'df'),
+    [
+        (
+            'one-dim',
+            parameterize.xrand(
+                (2,),
+                dtype='float32',
+                min=np.finfo(dtype='float32').tiny,
+            ),
+        ),
+        (
+            'multi-dim',
+            parameterize.xrand(
+                (2, 2),
+                dtype='float32',
+                min=np.finfo(dtype='float32').tiny,
+            ),
+        ),
+    ],
+)
+class TestChi2Sample(unittest.TestCase):
+    def setUp(self):
+        self.program = paddle.static.Program()
+        self.executor = paddle.static.Executor(self.place)
+        with paddle.static.program_guard(self.program):
+            df = paddle.static.data('df', self.df.shape, self.df.dtype)
+            self._paddle_chi2 = chi2.Chi2(df)
+            self.feeds = {'df': self.df}
+
+    def test_sample_shape(self):
+        cases = [
+            {
+                'input': (),
+                'expect': tuple(np.squeeze(self.df).shape),
+            },
+            {
+                'input': (2, 2),
+                'expect': (2, 2, *np.squeeze(self.df).shape),
+            },
+        ]
+        for case in cases:
+            with paddle.static.program_guard(self.program):
+                [data] = self.executor.run(
+                    self.program,
+                    feed=self.feeds,
+                    fetch_list=self._paddle_chi2.sample(case.get('input')),
+                )
+                self.assertTrue(data.shape == case.get('expect'))
+
+    def test_sample(self):
+        sample_shape = (30000,)
+        with paddle.static.program_guard(self.program):
+            [data] = self.executor.run(
+                self.program,
+                feed=self.feeds,
+                fetch_list=self._paddle_chi2.sample(sample_shape),
+            )
+            except_shape = sample_shape + np.squeeze(self.df).shape
+            self.assertTrue(data.shape == except_shape)
+            np.testing.assert_allclose(
+                data.mean(axis=0),
+                scipy.stats.chi2.mean(self.df),
+                rtol=0.1,
+                atol=config.ATOL.get(str(self.df.dtype)),
+            )
+            np.testing.assert_allclose(
+                data.var(axis=0),
+                scipy.stats.chi2.var(self.df),
+                rtol=0.1,
+                atol=config.ATOL.get(str(self.df.dtype)),
+            )
+
+
+@parameterize.place(config.DEVICES)
+@parameterize.parameterize_cls(
+    (parameterize.TEST_CASE_NAME, 'df'),
+    [
+        ('0-dim', 0.4),
+    ],
+)
+class TestChi2SampleKS(unittest.TestCase):
+    def setUp(self):
+        self.program = paddle.static.Program()
+        self.executor = paddle.static.Executor(self.place)
+        with paddle.static.program_guard(self.program):
+            df = paddle.static.data('df', (), 'float')
+            self._paddle_chi2 = chi2.Chi2(df)
+            self.feeds = {'df': self.df}
+
+    def test_sample(self):
+        sample_shape = (15000,)
+        with paddle.static.program_guard(self.program):
+            [samples] = self.executor.run(
+                self.program,
+                feed=self.feeds,
+                fetch_list=self._paddle_chi2.sample(sample_shape),
+            )
+            self.assertTrue(self._kstest(samples))
+
+    def _kstest(self, samples):
+        # Uses the Kolmogorov-Smirnov test for goodness of fit.
+        ks, _ = scipy.stats.kstest(samples, scipy.stats.chi2(self.df).cdf)
+        return ks < 0.02
+
+
+@parameterize.place(config.DEVICES)
+@parameterize.parameterize_cls(
+    (parameterize.TEST_CASE_NAME),
+    [
+        ('chi2_test_err'),
+    ],
+)
+class Chi2TestError(unittest.TestCase):
+    def setUp(self):
+        self.program = paddle.static.Program()
+
+    @parameterize.parameterize_func([(10,)])  # not sequence object sample shape
+    def test_bad_sample_shape(self, shape):
+        with paddle.static.program_guard(self.program):
+            _chi2 = chi2.Chi2(1.0)
+            self.assertRaises(TypeError, _chi2.sample, shape)
+
+
+if __name__ == '__main__':
+    unittest.main()

test/distribution/test_distribution_constraint.py

@@ -0,0 +1,89 @@
+# Copyright (c) 2022 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
+#
+#     http://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
+import parameterize as param
+
+import paddle
+from paddle.distribution import constraint
+
+np.random.seed(2022)
+
+
+@param.param_cls(
+    (param.TEST_CASE_NAME, 'value'), [('NotImplement', np.random.rand(2, 3))]
+)
+class TestConstraint(unittest.TestCase):
+    def setUp(self):
+        self._constraint = constraint.Constraint()
+
+    def test_costraint(self):
+        with self.assertRaises(NotImplementedError):
+            self._constraint.check(self.value)
+
+
+@param.param_cls(
+    (param.TEST_CASE_NAME, 'value', 'expect'), [('real', 1.0, True)]
+)
+class TestReal(unittest.TestCase):
+    def setUp(self):
+        self._constraint = constraint.Real()
+
+    def test_costraint(self):
+        self.assertEqual(self._constraint.check(self.value), self.expect)
+
+
+@param.param_cls(
+    (param.TEST_CASE_NAME, 'lower', 'upper', 'value', 'expect'),
+    [('in_range', 0, 1, 0.5, True), ('out_range', 0, 1, 2, False)],
+)
+class TestRange(unittest.TestCase):
+    def setUp(self):
+        self._constraint = constraint.Range(self.lower, self.upper)
+
+    def test_costraint(self):
+        self.assertEqual(self._constraint.check(self.value), self.expect)
+
+
+@param.param_cls(
+    (param.TEST_CASE_NAME, 'value', 'expect'),
+    [('positive', 1, True), ('negative', -1, False)],
+)
+class TestPositive(unittest.TestCase):
+    def setUp(self):
+        self._constraint = constraint.Positive()
+
+    def test_costraint(self):
+        self.assertEqual(self._constraint.check(self.value), self.expect)
+
+
+@param.param_cls(
+    (param.TEST_CASE_NAME, 'value', 'expect'),
+    [
+        ('simplex', paddle.to_tensor([0.5, 0.5]), True),
+        ('non_simplex', paddle.to_tensor([-0.5, 0.5]), False),
+    ],
+)
+class TestSimplex(unittest.TestCase):
+    def setUp(self):
+        self._constraint = constraint.Simplex()
+
+    def test_costraint(self):
+        self.assertEqual(self._constraint.check(self.value), self.expect)
+
+
+if __name__ == '__main__':
+    unittest.main()

test/distribution/test_distribution_continuous_bernoulli.py

@@ -0,0 +1,370 @@
+# Copyright (c) 2021 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
+#
+#     http://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
+import parameterize
+from distribution import config
+from parameterize import (
+    TEST_CASE_NAME,
+    parameterize_cls,
+    parameterize_func,
+)
+
+import paddle
+from paddle.distribution.continuous_bernoulli import ContinuousBernoulli
+
+
+class ContinuousBernoulli_np:
+    def __init__(self, probs, lims=(0.48, 0.52)):
+        self.lims = lims
+        self.dtype = probs.dtype
+        eps_prob = 1.1920928955078125e-07
+        self.probs = np.clip(probs, a_min=eps_prob, a_max=1.0 - eps_prob)
+
+    def _cut_support_region(self):
+        return np.logical_or(
+            np.less_equal(self.probs, self.lims[0]),
+            np.greater_equal(self.probs, self.lims[1]),
+        )
+
+    def _cut_probs(self):
+        return np.where(
+            self._cut_support_region(),
+            self.probs,
+            self.lims[0] * np.ones_like(self.probs),
+        )
+
+    def _tanh_inverse(self, value):
+        return 0.5 * (np.log1p(value) - np.log1p(-value))
+
+    def _log_constant(self):
+        cut_probs = self._cut_probs()
+        cut_probs_below_half = np.where(
+            np.less_equal(cut_probs, 0.5), cut_probs, np.zeros_like(cut_probs)
+        )
+        cut_probs_above_half = np.where(
+            np.greater_equal(cut_probs, 0.5), cut_probs, np.ones_like(cut_probs)
+        )
+        log_constant_propose = np.log(
+            2.0 * np.abs(self._tanh_inverse(1.0 - 2.0 * cut_probs))
+        ) - np.where(
+            np.less_equal(cut_probs, 0.5),
+            np.log1p(-2.0 * cut_probs_below_half),
+            np.log(2.0 * cut_probs_above_half - 1.0),
+        )
+        x = np.square(self.probs - 0.5)
+        taylor_expansion = np.log(2.0) + (4.0 / 3.0 + 104.0 / 45.0 * x) * x
+        return np.where(
+            self._cut_support_region(), log_constant_propose, taylor_expansion
+        )
+
+    def np_variance(self):
+        cut_probs = self._cut_probs()
+        tmp = np.divide(
+            np.square(cut_probs) - cut_probs, np.square(1.0 - 2.0 * cut_probs)
+        )
+        propose = tmp + np.divide(
+            1.0, np.square(2.0 * self._tanh_inverse(1.0 - 2.0 * cut_probs))
+        )
+        x = np.square(self.probs - 0.5)
+        taylor_expansion = 1.0 / 12.0 - (1.0 / 15.0 - 128.0 / 945.0 * x) * x
+        return np.where(self._cut_support_region(), propose, taylor_expansion)
+
+    def np_mean(self):
+        cut_probs = self._cut_probs()
+        tmp = cut_probs / (2.0 * cut_probs - 1.0)
+        propose = tmp + 1.0 / (2.0 * self._tanh_inverse(1.0 - 2.0 * cut_probs))
+        x = self.probs - 0.5
+        taylor_expansion = 0.5 + (1.0 / 3.0 + 16.0 / 45.0 * np.square(x)) * x
+        return np.where(self._cut_support_region(), propose, taylor_expansion)
+
+    def np_entropy(self):
+        log_p = np.log(self.probs)
+        log_1_minus_p = np.log1p(-self.probs)
+        return np.where(
+            np.equal(self.probs, 0.5),
+            np.full_like(self.probs, 0.0),
+            (
+                -self._log_constant()
+                + self.np_mean() * (log_1_minus_p - log_p)
+                - log_1_minus_p
+            ),
+        )
+
+    def np_prob(self, value):
+        return np.exp(self.np_log_prob(value))
+
+    def np_log_prob(self, value):
+        eps = 1e-8
+        cross_entropy = np.nan_to_num(
+            value * np.log(self.probs) + (1.0 - value) * np.log(1 - self.probs),
+            neginf=-eps,
+        )
+        return self._log_constant() + cross_entropy
+
+    def np_cdf(self, value):
+        cut_probs = self._cut_probs()
+        cdfs = (
+            np.power(cut_probs, value) * np.power(1.0 - cut_probs, 1.0 - value)
+            + cut_probs
+            - 1.0
+        ) / (2.0 * cut_probs - 1.0)
+        unbounded_cdfs = np.where(self._cut_support_region(), cdfs, value)
+        return np.where(
+            np.less_equal(value, 0.0),
+            np.zeros_like(value),
+            np.where(
+                np.greater_equal(value, 1.0),
+                np.ones_like(value),
+                unbounded_cdfs,
+            ),
+        )
+
+    def np_icdf(self, value):
+        cut_probs = self._cut_probs()
+        return np.where(
+            self._cut_support_region(),
+            (
+                np.log1p(-cut_probs + value * (2.0 * cut_probs - 1.0))
+                - np.log1p(-cut_probs)
+            )
+            / (np.log(cut_probs) - np.log1p(-cut_probs)),
+            value,
+        )
+
+    def np_kl_divergence(self, other):
+        part1 = -self.np_entropy()
+        log_q = np.log(other.probs)
+        log_1_minus_q = np.log1p(-other.probs)
+        part2 = -(
+            other._log_constant()
+            + self.np_mean() * (log_q - log_1_minus_q)
+            + log_1_minus_q
+        )
+        return part1 + part2
+
+
+@parameterize.place(config.DEVICES)
+@parameterize.parameterize_cls(
+    (parameterize.TEST_CASE_NAME, 'probs'),
+    [
+        ('half', np.array(0.5).astype("float32")),
+        (
+            'one-dim',
+            parameterize.xrand((1,), min=0.0, max=1.0).astype("float64"),
+        ),
+        (
+            'multi-dim',
+            parameterize.xrand((2, 3), min=0.0, max=1.0).astype("float32"),
+        ),
+    ],
+)
+class TestContinuousBernoulli(unittest.TestCase):
+    def setUp(self):
+        self._dist = ContinuousBernoulli(
+            probs=paddle.to_tensor(self.probs), lims=(0.48, 0.52)
+        )
+        self._np_dist = ContinuousBernoulli_np(self.probs, lims=(0.48, 0.52))
+
+    def test_mean(self):
+        mean = self._dist.mean
+        self.assertEqual(mean.numpy().dtype, self.probs.dtype)
+        np.testing.assert_allclose(
+            mean,
+            self._np_dist.np_mean(),
+            rtol=config.RTOL.get(str(self.probs.dtype)),
+            atol=config.ATOL.get(str(self.probs.dtype)),
+        )
+
+    def test_variance(self):
+        var = self._dist.variance
+        self.assertEqual(var.numpy().dtype, self.probs.dtype)
+        np.testing.assert_allclose(
+            var,
+            self._np_dist.np_variance(),
+            rtol=0.01,
+            atol=0.0,
+        )
+
+    def test_entropy(self):
+        entropy = self._dist.entropy()
+        self.assertEqual(entropy.numpy().dtype, self.probs.dtype)
+        np.testing.assert_allclose(
+            entropy,
+            self._np_dist.np_entropy(),
+            rtol=0.01,
+            atol=0.0,
+        )
+
+    def test_sample(self):
+        sample_shape = ()
+        samples = self._dist.sample(sample_shape)
+        self.assertEqual(samples.numpy().dtype, self.probs.dtype)
+        self.assertEqual(
+            tuple(samples.shape),
+            sample_shape + self._dist.batch_shape + self._dist.event_shape,
+        )
+
+        sample_shape = (50000,)
+        samples = self._dist.sample(sample_shape)
+        sample_mean = samples.mean(axis=0)
+        sample_variance = samples.var(axis=0)
+
+        np.testing.assert_allclose(
+            sample_mean,
+            self._dist.mean,
+            rtol=0.1,
+            atol=0.0,
+        )
+        np.testing.assert_allclose(
+            sample_variance,
+            self._dist.variance,
+            rtol=0.1,
+            atol=0.0,
+        )
+
+
+@parameterize.place(config.DEVICES)
+@parameterize.parameterize_cls(
+    (parameterize.TEST_CASE_NAME, 'probs', 'value'),
+    [
+        (
+            'zero-dim',
+            np.array(0.3).astype("float32"),
+            parameterize.xrand((5,), min=0.0, max=1.0).astype("float32"),
+        ),
+        (
+            'value-same-shape',
+            parameterize.xrand((5,), min=0.0, max=1.0).astype("float32"),
+            parameterize.xrand((5,), min=0.0, max=1.0).astype("float32"),
+        ),
+        (
+            'value-broadcast-shape',
+            parameterize.xrand((1,), min=0.0, max=1.0).astype("float64"),
+            parameterize.xrand((2, 3), min=0.0, max=1.0).astype("float64"),
+        ),
+    ],
+)
+class TestContinuousBernoulliProbs(unittest.TestCase):
+    def setUp(self):
+        self._dist = ContinuousBernoulli(
+            probs=paddle.to_tensor(self.probs), lims=(0.48, 0.52)
+        )
+        self._np_dist = ContinuousBernoulli_np(self.probs, lims=(0.48, 0.52))
+
+    def test_prob(self):
+        np.testing.assert_allclose(
+            self._dist.prob(paddle.to_tensor(self.value)),
+            self._np_dist.np_prob(self.value),
+            rtol=config.RTOL.get(str(self.probs.dtype)),
+            atol=config.ATOL.get(str(self.probs.dtype)),
+        )
+
+    def test_log_prob(self):
+        np.testing.assert_allclose(
+            self._dist.log_prob(paddle.to_tensor(self.value)),
+            self._np_dist.np_log_prob(self.value),
+            rtol=config.RTOL.get(str(self.probs.dtype)),
+            atol=config.ATOL.get(str(self.probs.dtype)),
+        )
+
+    def test_cdf(self):
+        np.testing.assert_allclose(
+            self._dist.cdf(paddle.to_tensor(self.value)),
+            self._np_dist.np_cdf(self.value),
+            rtol=config.RTOL.get(str(self.probs.dtype)),
+            atol=config.ATOL.get(str(self.probs.dtype)),
+        )
+
+    def test_icdf(self):
+        np.testing.assert_allclose(
+            self._dist.icdf(paddle.to_tensor(self.value)),
+            self._np_dist.np_icdf(self.value),
+            rtol=config.RTOL.get(str(self.probs.dtype)),
+            atol=config.ATOL.get(str(self.probs.dtype)),
+        )
+
+
+@parameterize.place(config.DEVICES)
+@parameterize.parameterize_cls(
+    (parameterize.TEST_CASE_NAME, 'p_1', 'p_2'),
+    [
+        (
+            'zero-dim',
+            np.array(0.2).astype("float32"),
+            np.array(0.4).astype("float32"),
+        ),
+        (
+            'one-dim',
+            parameterize.xrand((1,), min=0.0, max=1.0).astype("float32"),
+            parameterize.xrand((1,), min=0.0, max=1.0).astype("float32"),
+        ),
+        (
+            'multi-dim',
+            parameterize.xrand((5,), min=0.0, max=1.0).astype("float64"),
+            parameterize.xrand((5,), min=0.0, max=1.0).astype("float64"),
+        ),
+    ],
+)
+class TestContinuousBernoulliKL(unittest.TestCase):
+    def setUp(self):
+        paddle.disable_static()
+        self._dist1 = ContinuousBernoulli(
+            probs=paddle.to_tensor(self.p_1), lims=(0.48, 0.52)
+        )
+        self._dist2 = ContinuousBernoulli(
+            probs=paddle.to_tensor(self.p_2), lims=(0.48, 0.52)
+        )
+        self._np_dist1 = ContinuousBernoulli_np(self.p_1, lims=(0.48, 0.52))
+        self._np_dist2 = ContinuousBernoulli_np(self.p_2, lims=(0.48, 0.52))
+
+    def test_kl_divergence(self):
+        kl0 = self._dist1.kl_divergence(self._dist2)
+        kl1 = self._np_dist1.np_kl_divergence(self._np_dist2)
+
+        self.assertEqual(tuple(kl0.shape), self._dist1.batch_shape)
+        self.assertEqual(tuple(kl1.shape), self._dist1.batch_shape)
+        np.testing.assert_allclose(
+            kl0,
+            kl1,
+            rtol=0.01,
+            atol=0.0,
+        )
+
+
+@parameterize.place(config.DEVICES)
+@parameterize_cls([TEST_CASE_NAME], ['ContinuousBernoulliTestError'])
+class ContinuousBernoulliTestError(unittest.TestCase):
+    def setUp(self):
+        paddle.disable_static(self.place)
+
+    @parameterize_func(
+        [
+            (
+                paddle.to_tensor([0.3, 0.5]),
+                paddle.to_tensor([0.2, 0.8, 0.6]),
+            ),
+        ]
+    )
+    def test_bad_kl_div(self, probs1, probs2):
+        with paddle.base.dygraph.guard(self.place):
+            rv = ContinuousBernoulli(probs1)
+            rv_other = ContinuousBernoulli(probs2)
+            self.assertRaises(ValueError, rv.kl_divergence, rv_other)
+
+
+if __name__ == '__main__':
+    unittest.main(argv=[''], verbosity=3, exit=False)

test/distribution/test_distribution_continuous_bernoulli_static.py

@@ -0,0 +1,336 @@
+# Copyright (c) 2021 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
+#
+#     http://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
+import parameterize
+from distribution import config
+
+import paddle
+from paddle.distribution.continuous_bernoulli import ContinuousBernoulli
+
+
+class ContinuousBernoulli_np:
+    def __init__(self, probs, lims=(0.48, 0.52)):
+        self.lims = lims
+        self.dtype = probs.dtype
+        eps_prob = 1.1920928955078125e-07
+        self.probs = np.clip(probs, a_min=eps_prob, a_max=1.0 - eps_prob)
+
+    def _cut_support_region(self):
+        return np.logical_or(
+            np.less_equal(self.probs, self.lims[0]),
+            np.greater_equal(self.probs, self.lims[1]),
+        )
+
+    def _cut_probs(self):
+        return np.where(
+            self._cut_support_region(),
+            self.probs,
+            self.lims[0] * np.ones_like(self.probs),
+        )
+
+    def _tanh_inverse(self, value):
+        return 0.5 * (np.log1p(value) - np.log1p(-value))
+
+    def _log_constant(self):
+        cut_probs = self._cut_probs()
+        cut_probs_below_half = np.where(
+            np.less_equal(cut_probs, 0.5), cut_probs, np.zeros_like(cut_probs)
+        )
+        cut_probs_above_half = np.where(
+            np.greater_equal(cut_probs, 0.5), cut_probs, np.ones_like(cut_probs)
+        )
+        log_constant_propose = np.log(
+            2.0 * np.abs(self._tanh_inverse(1.0 - 2.0 * cut_probs))
+        ) - np.where(
+            np.less_equal(cut_probs, 0.5),
+            np.log1p(-2.0 * cut_probs_below_half),
+            np.log(2.0 * cut_probs_above_half - 1.0),
+        )
+        x = np.square(self.probs - 0.5)
+        taylor_expansion = np.log(2.0) + (4.0 / 3.0 + 104.0 / 45.0 * x) * x
+        return np.where(
+            self._cut_support_region(), log_constant_propose, taylor_expansion
+        )
+
+    def np_variance(self):
+        cut_probs = self._cut_probs()
+        tmp = np.divide(
+            np.square(cut_probs) - cut_probs, np.square(1.0 - 2.0 * cut_probs)
+        )
+        propose = tmp + np.divide(
+            1.0, np.square(2.0 * self._tanh_inverse(1.0 - 2.0 * cut_probs))
+        )
+        x = np.square(self.probs - 0.5)
+        taylor_expansion = 1.0 / 12.0 - (1.0 / 15.0 - 128.0 / 945.0 * x) * x
+        return np.where(self._cut_support_region(), propose, taylor_expansion)
+
+    def np_mean(self):
+        cut_probs = self._cut_probs()
+        tmp = cut_probs / (2.0 * cut_probs - 1.0)
+        propose = tmp + 1.0 / (2.0 * self._tanh_inverse(1.0 - 2.0 * cut_probs))
+        x = self.probs - 0.5
+        taylor_expansion = 0.5 + (1.0 / 3.0 + 16.0 / 45.0 * np.square(x)) * x
+        return np.where(self._cut_support_region(), propose, taylor_expansion)
+
+    def np_entropy(self):
+        log_p = np.log(self.probs)
+        log_1_minus_p = np.log1p(-self.probs)
+        return np.where(
+            np.equal(self.probs, 0.5),
+            np.full_like(self.probs, 0.0),
+            (
+                -self._log_constant()
+                + self.np_mean() * (log_1_minus_p - log_p)
+                - log_1_minus_p
+            ),
+        )
+
+    def np_prob(self, value):
+        return np.exp(self.np_log_prob(value))
+
+    def np_log_prob(self, value):
+        eps = 1e-8
+        cross_entropy = np.nan_to_num(
+            value * np.log(self.probs) + (1.0 - value) * np.log(1 - self.probs),
+            neginf=-eps,
+        )
+        return self._log_constant() + cross_entropy
+
+    def np_cdf(self, value):
+        cut_probs = self._cut_probs()
+        cdfs = (
+            np.power(cut_probs, value) * np.power(1.0 - cut_probs, 1.0 - value)
+            + cut_probs
+            - 1.0
+        ) / (2.0 * cut_probs - 1.0)
+        unbounded_cdfs = np.where(self._cut_support_region(), cdfs, value)
+        return np.where(
+            np.less_equal(value, 0.0),
+            np.zeros_like(value),
+            np.where(
+                np.greater_equal(value, 1.0),
+                np.ones_like(value),
+                unbounded_cdfs,
+            ),
+        )
+
+    def np_icdf(self, value):
+        cut_probs = self._cut_probs()
+        return np.where(
+            self._cut_support_region(),
+            (
+                np.log1p(-cut_probs + value * (2.0 * cut_probs - 1.0))
+                - np.log1p(-cut_probs)
+            )
+            / (np.log(cut_probs) - np.log1p(-cut_probs)),
+            value,
+        )
+
+    def np_kl_divergence(self, other):
+        part1 = -self.np_entropy()
+        log_q = np.log(other.probs)
+        log_1_minus_q = np.log1p(-other.probs)
+        part2 = -(
+            other._log_constant()
+            + self.np_mean() * (log_q - log_1_minus_q)
+            + log_1_minus_q
+        )
+        return part1 + part2
+
+
+paddle.enable_static()
+
+
+@parameterize.place(config.DEVICES)
+@parameterize.parameterize_cls(
+    (parameterize.TEST_CASE_NAME, 'probs'),
+    [
+        (
+            'zero-dim',
+            np.array(0.7).astype("float32"),
+        ),
+        (
+            'multi-dim',
+            parameterize.xrand((1, 3), min=0.0, max=1.0).astype("float32"),
+        ),
+    ],
+)
+class TestContinuousBernoulli(unittest.TestCase):
+    def setUp(self):
+        self._np_dist = ContinuousBernoulli_np(self.probs)
+        startup_program = paddle.static.Program()
+        main_program = paddle.static.Program()
+        executor = paddle.static.Executor(self.place)
+        with paddle.static.program_guard(main_program, startup_program):
+            probs = paddle.static.data(
+                'probs', self.probs.shape, self.probs.dtype
+            )
+            dist = ContinuousBernoulli(probs, lims=(0.48, 0.52))
+            mean = dist.mean
+            var = dist.variance
+            entropy = dist.entropy()
+            large_samples = dist.sample(shape=(50000,))
+        fetch_list = [mean, var, entropy, large_samples]
+        feed = {'probs': self.probs}
+
+        executor.run(startup_program)
+        [
+            self.mean,
+            self.var,
+            self.entropy,
+            self.large_samples,
+        ] = executor.run(main_program, feed=feed, fetch_list=fetch_list)
+
+    def test_mean(self):
+        self.assertEqual(str(self.mean.dtype).split('.')[-1], self.probs.dtype)
+        np.testing.assert_allclose(
+            self.mean,
+            self._np_mean(),
+            rtol=config.RTOL.get(str(self.probs.dtype)),
+            atol=config.ATOL.get(str(self.probs.dtype)),
+        )
+
+    def test_variance(self):
+        self.assertEqual(str(self.var.dtype).split('.')[-1], self.probs.dtype)
+        np.testing.assert_allclose(
+            self.var,
+            self._np_variance(),
+            rtol=0.01,
+            atol=0.0,
+        )
+
+    def test_entropy(self):
+        self.assertEqual(
+            str(self.entropy.dtype).split('.')[-1], self.probs.dtype
+        )
+        np.testing.assert_allclose(
+            self.entropy,
+            self._np_entropy(),
+            rtol=0.01,
+            atol=0.0,
+        )
+
+    def test_sample(self):
+        sample_mean = self.large_samples.mean(axis=0)
+        sample_variance = self.large_samples.var(axis=0)
+        np.testing.assert_allclose(sample_mean, self.mean, atol=0, rtol=0.1)
+        np.testing.assert_allclose(sample_variance, self.var, atol=0, rtol=0.1)
+
+    def _np_variance(self):
+        return self._np_dist.np_variance()
+
+    def _np_mean(self):
+        return self._np_dist.np_mean()
+
+    def _np_entropy(self):
+        return self._np_dist.np_entropy()
+
+
+@parameterize.place(config.DEVICES)
+@parameterize.parameterize_cls(
+    (parameterize.TEST_CASE_NAME, 'probs', 'value'),
+    [
+        (
+            'value-broadcast-shape',
+            parameterize.xrand((1,), min=0.0, max=1.0).astype("float32"),
+            parameterize.xrand((2, 2), min=0.0, max=1.0).astype("float64"),
+        ),
+    ],
+)
+class TestContinuousBernoulliProbs(unittest.TestCase):
+    def setUp(self):
+        self._np_dist = ContinuousBernoulli_np(self.probs)
+        startup_program = paddle.static.Program()
+        main_program = paddle.static.Program()
+        executor = paddle.static.Executor(self.place)
+
+        with paddle.static.program_guard(main_program, startup_program):
+            probs = paddle.static.data(
+                'probs', self.probs.shape, self.probs.dtype
+            )
+            value = paddle.static.data(
+                'value', self.value.shape, self.value.dtype
+            )
+            dist = ContinuousBernoulli(probs, lims=(0.48, 0.52))
+            pmf = dist.prob(value)
+        feed = {'probs': self.probs, 'value': self.value}
+        fetch_list = [pmf]
+
+        executor.run(startup_program)
+        [self.pmf] = executor.run(
+            main_program, feed=feed, fetch_list=fetch_list
+        )
+
+    def test_prob(self):
+        np.testing.assert_allclose(
+            self.pmf,
+            self._np_dist.np_prob(self.value),
+            rtol=config.RTOL.get(str(self.probs.dtype)),
+            atol=config.ATOL.get(str(self.probs.dtype)),
+        )
+
+
+@parameterize.place(config.DEVICES)
+@parameterize.parameterize_cls(
+    (parameterize.TEST_CASE_NAME, 'p_1', 'p_2'),
+    [
+        (
+            'multi-dim',
+            parameterize.xrand((2,), min=0.0, max=1.0).astype("float32"),
+            parameterize.xrand((2,), min=0.0, max=1.0).astype("float32"),
+        ),
+    ],
+)
+class TestContinuousBernoulliKL(unittest.TestCase):
+    def setUp(self):
+        self._np_dist1 = ContinuousBernoulli_np(self.p_1)
+        self._np_dist2 = ContinuousBernoulli_np(self.p_2)
+
+        startup_program = paddle.static.Program()
+        main_program = paddle.static.Program()
+        executor = paddle.static.Executor(self.place)
+
+        with paddle.static.program_guard(main_program, startup_program):
+            p_1 = paddle.static.data('p_1', self.p_1.shape)
+            p_2 = paddle.static.data('p_2', self.p_2.shape)
+            dist1 = ContinuousBernoulli(p_1, lims=(0.48, 0.52))
+            dist2 = ContinuousBernoulli(p_2, lims=(0.48, 0.52))
+            kl_dist1_dist2 = dist1.kl_divergence(dist2)
+        feed = {'p_1': self.p_1, 'p_2': self.p_2}
+        fetch_list = [kl_dist1_dist2]
+
+        executor.run(startup_program)
+        [self.kl_dist1_dist2] = executor.run(
+            main_program, feed=feed, fetch_list=fetch_list
+        )
+
+    def test_kl_divergence(self):
+        kl0 = self.kl_dist1_dist2
+        kl1 = self._np_dist1.np_kl_divergence(self._np_dist2)
+
+        self.assertEqual(tuple(kl0.shape), self.p_1.shape)
+        self.assertEqual(tuple(kl1.shape), self.p_1.shape)
+        np.testing.assert_allclose(
+            kl0,
+            kl1,
+            rtol=0.01,
+            atol=0.0,
+        )
+
+
+if __name__ == '__main__':
+    unittest.main(argv=[''], verbosity=3, exit=False)

test/distribution/test_distribution_dirichlet.py

@@ -0,0 +1,123 @@
+# Copyright (c) 2021 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
+#
+#     http://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
+import parameterize as param
+import scipy.stats
+from distribution.config import ATOL, DEVICES, RTOL
+
+import paddle
+
+np.random.seed(2022)
+
+
+@param.place(DEVICES)
+@param.param_cls(
+    (param.TEST_CASE_NAME, 'concentration'),
+    [
+        ('test-one-dim', param.xrand((89,))),
+        # ('test-multi-dim', config.xrand((10, 20, 30)))
+    ],
+)
+class TestDirichlet(unittest.TestCase):
+    def setUp(self):
+        self._paddle_diric = paddle.distribution.Dirichlet(
+            paddle.to_tensor(self.concentration)
+        )
+
+    def test_mean(self):
+        with paddle.base.dygraph.guard(self.place):
+            np.testing.assert_allclose(
+                self._paddle_diric.mean,
+                scipy.stats.dirichlet.mean(self.concentration),
+                rtol=RTOL.get(str(self.concentration.dtype)),
+                atol=ATOL.get(str(self.concentration.dtype)),
+            )
+
+    def test_variance(self):
+        with paddle.base.dygraph.guard(self.place):
+            np.testing.assert_allclose(
+                self._paddle_diric.variance,
+                scipy.stats.dirichlet.var(self.concentration),
+                rtol=RTOL.get(str(self.concentration.dtype)),
+                atol=ATOL.get(str(self.concentration.dtype)),
+            )
+
+    def test_prob(self):
+        value = [np.random.rand(*self.concentration.shape)]
+        value = [v / v.sum() for v in value]
+
+        for v in value:
+            with paddle.base.dygraph.guard(self.place):
+                np.testing.assert_allclose(
+                    self._paddle_diric.prob(paddle.to_tensor(v)),
+                    scipy.stats.dirichlet.pdf(v, self.concentration),
+                    rtol=RTOL.get(str(self.concentration.dtype)),
+                    atol=ATOL.get(str(self.concentration.dtype)),
+                )
+
+    def test_log_prob(self):
+        value = [np.random.rand(*self.concentration.shape)]
+        value = [v / v.sum() for v in value]
+
+        for v in value:
+            with paddle.base.dygraph.guard(self.place):
+                np.testing.assert_allclose(
+                    self._paddle_diric.log_prob(paddle.to_tensor(v)),
+                    scipy.stats.dirichlet.logpdf(v, self.concentration),
+                    rtol=RTOL.get(str(self.concentration.dtype)),
+                    atol=ATOL.get(str(self.concentration.dtype)),
+                )
+
+    def test_entropy(self):
+        with paddle.base.dygraph.guard(self.place):
+            np.testing.assert_allclose(
+                self._paddle_diric.entropy(),
+                scipy.stats.dirichlet.entropy(self.concentration),
+                rtol=RTOL.get(str(self.concentration.dtype)),
+                atol=ATOL.get(str(self.concentration.dtype)),
+            )
+
+    def test_natural_parameters(self):
+        self.assertTrue(
+            isinstance(self._paddle_diric._natural_parameters, tuple)
+        )
+
+    def test_log_normalizer(self):
+        self.assertTrue(
+            np.all(
+                self._paddle_diric._log_normalizer(
+                    paddle.to_tensor(param.xrand((100, 100, 100)))
+                ).numpy()
+                < 0.0
+            )
+        )
+
+    @param.place(DEVICES)
+    @param.param_cls(
+        (param.TEST_CASE_NAME, 'concentration'),
+        [('test-zero-dim', np.array(1.0))],
+    )
+    class TestDirichletException(unittest.TestCase):
+        def TestInit(self):
+            with self.assertRaises(ValueError):
+                paddle.distribution.Dirichlet(
+                    paddle.squeeze(self.concentration)
+                )
+
+
+if __name__ == '__main__':
+    unittest.main()

test/distribution/test_distribution_dirichlet_static.py

@@ -0,0 +1,122 @@
+# Copyright (c) 2021 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
+#
+#     http://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
+import scipy.stats
+from distribution.config import ATOL, DEVICES, RTOL
+from parameterize import TEST_CASE_NAME, parameterize_cls, place
+
+import paddle
+
+np.random.seed(2022)
+paddle.enable_static()
+
+
+@place(DEVICES)
+@parameterize_cls(
+    (TEST_CASE_NAME, 'concentration'),
+    [('test-one-dim', np.random.rand(89) + 5.0)],
+)
+class TestDirichlet(unittest.TestCase):
+    def setUp(self):
+        self.program = paddle.static.Program()
+        self.executor = paddle.static.Executor()
+        with paddle.static.program_guard(self.program):
+            conc = paddle.static.data(
+                'conc', self.concentration.shape, self.concentration.dtype
+            )
+            self._paddle_diric = paddle.distribution.Dirichlet(conc)
+            self.feeds = {'conc': self.concentration}
+
+    def test_mean(self):
+        with paddle.static.program_guard(self.program):
+            [out] = self.executor.run(
+                self.program,
+                feed=self.feeds,
+                fetch_list=[self._paddle_diric.mean],
+            )
+            np.testing.assert_allclose(
+                out,
+                scipy.stats.dirichlet.mean(self.concentration),
+                rtol=RTOL.get(str(self.concentration.dtype)),
+                atol=ATOL.get(str(self.concentration.dtype)),
+            )
+
+    def test_variance(self):
+        with paddle.static.program_guard(self.program):
+            [out] = self.executor.run(
+                self.program,
+                feed=self.feeds,
+                fetch_list=[self._paddle_diric.variance],
+            )
+            np.testing.assert_allclose(
+                out,
+                scipy.stats.dirichlet.var(self.concentration),
+                rtol=RTOL.get(str(self.concentration.dtype)),
+                atol=ATOL.get(str(self.concentration.dtype)),
+            )
+
+    def test_prob(self):
+        with paddle.static.program_guard(self.program):
+            random_number = np.random.rand(*self.concentration.shape)
+            random_number = random_number / random_number.sum()
+            feeds = dict(self.feeds, value=random_number)
+            value = paddle.static.data(
+                'value', random_number.shape, random_number.dtype
+            )
+            out = self._paddle_diric.prob(value)
+            [out] = self.executor.run(
+                self.program, feed=feeds, fetch_list=[out]
+            )
+            np.testing.assert_allclose(
+                out,
+                scipy.stats.dirichlet.pdf(random_number, self.concentration),
+                rtol=RTOL.get(str(self.concentration.dtype)),
+                atol=ATOL.get(str(self.concentration.dtype)),
+            )
+
+    def test_log_prob(self):
+        with paddle.static.program_guard(self.program):
+            random_number = np.random.rand(*self.concentration.shape)
+            random_number = random_number / random_number.sum()
+            feeds = dict(self.feeds, value=random_number)
+            value = paddle.static.data(
+                'value', random_number.shape, random_number.dtype
+            )
+            out = self._paddle_diric.log_prob(value)
+            [out] = self.executor.run(
+                self.program, feed=feeds, fetch_list=[out]
+            )
+            np.testing.assert_allclose(
+                out,
+                scipy.stats.dirichlet.logpdf(random_number, self.concentration),
+                rtol=RTOL.get(str(self.concentration.dtype)),
+                atol=ATOL.get(str(self.concentration.dtype)),
+            )
+
+    def test_entropy(self):
+        with paddle.static.program_guard(self.program):
+            [out] = self.executor.run(
+                self.program,
+                feed=self.feeds,
+                fetch_list=[self._paddle_diric.entropy()],
+            )
+            np.testing.assert_allclose(
+                out,
+                scipy.stats.dirichlet.entropy(self.concentration),
+                rtol=RTOL.get(str(self.concentration.dtype)),
+                atol=ATOL.get(str(self.concentration.dtype)),
+            )

test/distribution/test_distribution_expfamily.py

@@ -0,0 +1,76 @@
+# Copyright (c) 2021 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
+#
+#     http://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 mock_data as mock
+import numpy as np
+import parameterize
+from distribution import config
+
+import paddle
+
+np.random.seed(2022)
+
+
+@parameterize.place(config.DEVICES)
+@parameterize.parameterize_cls(
+    (parameterize.TEST_CASE_NAME, 'dist'),
+    [
+        (
+            'test-mock-exp',
+            mock.Exponential(
+                rate=paddle.rand([100, 200, 99], dtype=config.DEFAULT_DTYPE)
+            ),
+        )
+    ],
+)
+class TestExponentialFamily(unittest.TestCase):
+    def test_entropy(self):
+        np.testing.assert_allclose(
+            self.dist.entropy(),
+            paddle.distribution.ExponentialFamily.entropy(self.dist),
+            rtol=config.RTOL.get(config.DEFAULT_DTYPE),
+            atol=config.ATOL.get(config.DEFAULT_DTYPE),
+        )
+
+
+@parameterize.place(config.DEVICES)
+@parameterize.parameterize_cls(
+    (config.TEST_CASE_NAME, 'dist'),
+    [
+        ('test-dummy', mock.DummyExpFamily(0.5, 0.5)),
+        (
+            'test-dirichlet',
+            paddle.distribution.Dirichlet(
+                paddle.to_tensor(parameterize.xrand())
+            ),
+        ),
+        (
+            'test-beta',
+            paddle.distribution.Beta(
+                paddle.to_tensor(parameterize.xrand()),
+                paddle.to_tensor(parameterize.xrand()),
+            ),
+        ),
+    ],
+)
+class TestExponentialFamilyException(unittest.TestCase):
+    def test_entropy_exception(self):
+        with self.assertRaises(NotImplementedError):
+            paddle.distribution.ExponentialFamily.entropy(self.dist)
+
+
+if __name__ == '__main__':
+    unittest.main()

test/distribution/test_distribution_expfamily_static.py

@@ -0,0 +1,70 @@
+# Copyright (c) 2021 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
+#
+#     http://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 mock_data as mock
+import numpy as np
+import parameterize
+from distribution import config
+
+import paddle
+
+np.random.seed(2022)
+paddle.enable_static()
+
+
+@parameterize.place(config.DEVICES)
+class TestExponentialFamily(unittest.TestCase):
+    def setUp(self):
+        self.program = paddle.static.Program()
+        self.executor = paddle.static.Executor()
+        with paddle.static.program_guard(self.program):
+            rate_np = parameterize.xrand((100, 200, 99))
+            rate = paddle.static.data('rate', rate_np.shape, rate_np.dtype)
+            self.mock_dist = mock.Exponential(rate)
+            self.feeds = {'rate': rate_np}
+
+    def test_entropy(self):
+        with paddle.static.program_guard(self.program):
+            [out1, out2] = self.executor.run(
+                self.program,
+                feed=self.feeds,
+                fetch_list=[
+                    self.mock_dist.entropy(),
+                    paddle.distribution.ExponentialFamily.entropy(
+                        self.mock_dist
+                    ),
+                ],
+            )
+
+            np.testing.assert_allclose(
+                out1,
+                out2,
+                rtol=config.RTOL.get(config.DEFAULT_DTYPE),
+                atol=config.ATOL.get(config.DEFAULT_DTYPE),
+            )
+
+    def test_entropy_exception(self):
+        with (
+            paddle.static.program_guard(self.program),
+            self.assertRaises(NotImplementedError),
+        ):
+            paddle.distribution.ExponentialFamily.entropy(
+                mock.DummyExpFamily(0.5, 0.5)
+            )
+
+
+if __name__ == '__main__':
+    unittest.main()

test/distribution/test_distribution_exponential.py

@@ -0,0 +1,370 @@
+# 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
+#
+#     http://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 numbers
+import unittest
+
+import numpy as np
+import parameterize
+import scipy.stats
+from distribution import config
+
+import paddle
+from paddle.distribution import exponential, kl
+
+np.random.seed(2023)
+paddle.seed(2023)
+
+
+@parameterize.place(config.DEVICES)
+@parameterize.parameterize_cls(
+    (parameterize.TEST_CASE_NAME, 'rate'),
+    [
+        (
+            '0-dim',
+            0.5,
+        ),
+        (
+            'one-dim',
+            parameterize.xrand(
+                (4,),
+                dtype='float32',
+                min=np.finfo(dtype='float32').tiny,
+            ),
+        ),
+        (
+            'multi-dim',
+            parameterize.xrand(
+                (10, 12),
+                dtype='float32',
+                min=np.finfo(dtype='float32').tiny,
+            ),
+        ),
+    ],
+)
+class TestExponential(unittest.TestCase):
+    def setUp(self):
+        rate = self.rate
+        if not isinstance(self.rate, numbers.Real):
+            rate = paddle.to_tensor(self.rate, dtype=paddle.float32)
+
+        self.scale = 1 / rate
+        self._paddle_expon = exponential.Exponential(rate)
+
+    def test_mean(self):
+        with paddle.base.dygraph.guard(self.place):
+            np.testing.assert_allclose(
+                self._paddle_expon.mean,
+                scipy.stats.expon.mean(scale=self.scale),
+                rtol=config.RTOL.get(
+                    str(self._paddle_expon.rate.numpy().dtype)
+                ),
+                atol=config.ATOL.get(
+                    str(self._paddle_expon.rate.numpy().dtype)
+                ),
+            )
+
+    def test_variance(self):
+        with paddle.base.dygraph.guard(self.place):
+            np.testing.assert_allclose(
+                self._paddle_expon.variance,
+                scipy.stats.expon.var(scale=self.scale),
+                rtol=config.RTOL.get(
+                    str(self._paddle_expon.rate.numpy().dtype)
+                ),
+                atol=config.ATOL.get(
+                    str(self._paddle_expon.rate.numpy().dtype)
+                ),
+            )
+
+    def test_prob(self):
+        value = np.random.rand(*self._paddle_expon.rate.shape)
+        with paddle.base.dygraph.guard(self.place):
+            np.testing.assert_allclose(
+                self._paddle_expon.prob(paddle.to_tensor(value)),
+                scipy.stats.expon.pdf(value, scale=self.scale),
+                rtol=config.RTOL.get(
+                    str(self._paddle_expon.rate.numpy().dtype)
+                ),
+                atol=config.ATOL.get(
+                    str(self._paddle_expon.rate.numpy().dtype)
+                ),
+            )
+
+    def test_cdf(self):
+        value = np.random.rand(*self._paddle_expon.rate.shape)
+        with paddle.base.dygraph.guard(self.place):
+            np.testing.assert_allclose(
+                self._paddle_expon.cdf(paddle.to_tensor(value)),
+                scipy.stats.expon.cdf(value, scale=self.scale),
+                rtol=config.RTOL.get(
+                    str(self._paddle_expon.rate.numpy().dtype)
+                ),
+                atol=config.ATOL.get(
+                    str(self._paddle_expon.rate.numpy().dtype)
+                ),
+            )
+
+    def test_icdf(self):
+        value = np.random.rand(*self._paddle_expon.rate.shape)
+        with paddle.base.dygraph.guard(self.place):
+            np.testing.assert_allclose(
+                self._paddle_expon.icdf(paddle.to_tensor(value)),
+                scipy.stats.expon.ppf(value, scale=self.scale),
+                rtol=config.RTOL.get(
+                    str(self._paddle_expon.rate.numpy().dtype)
+                ),
+                atol=config.ATOL.get(
+                    str(self._paddle_expon.rate.numpy().dtype)
+                ),
+            )
+
+    def test_entropy(self):
+        with paddle.base.dygraph.guard(self.place):
+            np.testing.assert_allclose(
+                self._paddle_expon.entropy(),
+                scipy.stats.expon.entropy(scale=self.scale),
+                rtol=config.RTOL.get(
+                    str(self._paddle_expon.rate.numpy().dtype)
+                ),
+                atol=config.ATOL.get(
+                    str(self._paddle_expon.rate.numpy().dtype)
+                ),
+            )
+
+
+@parameterize.place(config.DEVICES)
+@parameterize.parameterize_cls(
+    (parameterize.TEST_CASE_NAME, 'rate'),
+    [
+        (
+            'one-dim',
+            parameterize.xrand(
+                (2,),
+                dtype='float32',
+                min=np.finfo(dtype='float32').tiny,
+            ),
+        ),
+        (
+            'multi-dim',
+            parameterize.xrand(
+                (2, 3),
+                dtype='float32',
+                min=np.finfo(dtype='float32').tiny,
+            ),
+        ),
+    ],
+)
+class TestExponentialSample(unittest.TestCase):
+    def setUp(self):
+        rate = self.rate
+        if not isinstance(self.rate, numbers.Real):
+            rate = paddle.to_tensor(self.rate, dtype=paddle.float32)
+
+        self.scale = 1 / rate
+        self._paddle_expon = exponential.Exponential(rate)
+
+    def test_sample_shape(self):
+        cases = [
+            {
+                'input': (),
+                'expect': tuple(paddle.squeeze(self._paddle_expon.rate).shape),
+            },
+            {
+                'input': (3, 2),
+                'expect': (
+                    3,
+                    2,
+                    *paddle.squeeze(self._paddle_expon.rate).shape,
+                ),
+            },
+        ]
+        for case in cases:
+            self.assertTrue(
+                tuple(self._paddle_expon.sample(case.get('input')).shape)
+                == case.get('expect')
+            )
+
+    def test_sample(self):
+        sample_shape = (20000,)
+        samples = self._paddle_expon.sample(sample_shape)
+        sample_values = samples.numpy()
+        self.assertEqual(sample_values.dtype, self.rate.dtype)
+
+        np.testing.assert_allclose(
+            sample_values.mean(axis=0),
+            scipy.stats.expon.mean(scale=self.scale),
+            rtol=0.1,
+            atol=config.ATOL.get(str(self._paddle_expon.rate.numpy().dtype)),
+        )
+        np.testing.assert_allclose(
+            sample_values.var(axis=0),
+            scipy.stats.expon.var(scale=self.scale),
+            rtol=0.1,
+            atol=config.ATOL.get(str(self._paddle_expon.rate.numpy().dtype)),
+        )
+
+    def test_rsample_shape(self):
+        cases = [
+            {
+                'input': (),
+                'expect': tuple(paddle.squeeze(self._paddle_expon.rate).shape),
+            },
+            {
+                'input': (2, 5),
+                'expect': (
+                    2,
+                    5,
+                    *paddle.squeeze(self._paddle_expon.rate).shape,
+                ),
+            },
+        ]
+        for case in cases:
+            self.assertTrue(
+                tuple(self._paddle_expon.rsample(case.get('input')).shape)
+                == case.get('expect')
+            )
+
+    def test_rsample(self):
+        sample_shape = (20000,)
+        samples = self._paddle_expon.rsample(sample_shape)
+        sample_values = samples.numpy()
+        self.assertEqual(sample_values.dtype, self.rate.dtype)
+
+        np.testing.assert_allclose(
+            sample_values.mean(axis=0),
+            scipy.stats.expon.mean(scale=self.scale),
+            rtol=0.1,
+            atol=config.ATOL.get(str(self._paddle_expon.rate.numpy().dtype)),
+        )
+        np.testing.assert_allclose(
+            sample_values.var(axis=0),
+            scipy.stats.expon.var(scale=self.scale),
+            rtol=0.1,
+            atol=config.ATOL.get(str(self._paddle_expon.rate.numpy().dtype)),
+        )
+
+    def test_rsample_backpropagation(self):
+        sample_shape = (1000, 2)
+        with paddle.base.dygraph.guard(self.place):
+            self._paddle_expon.rate.stop_gradient = False
+            samples = self._paddle_expon.rsample(sample_shape)
+            grads = paddle.grad([samples], [self._paddle_expon.rate])
+            self.assertEqual(len(grads), 1)
+            self.assertEqual(grads[0].dtype, self._paddle_expon.rate.dtype)
+            self.assertEqual(grads[0].shape, self._paddle_expon.rate.shape)
+            axis = list(range(len(sample_shape)))
+            np.testing.assert_allclose(
+                -samples.sum(axis) / self._paddle_expon.rate,
+                grads[0],
+                rtol=config.RTOL.get(
+                    str(self._paddle_expon.rate.numpy().dtype)
+                ),
+                atol=config.ATOL.get(
+                    str(self._paddle_expon.rate.numpy().dtype)
+                ),
+            )
+
+
+@parameterize.place(config.DEVICES)
+@parameterize.parameterize_cls(
+    (parameterize.TEST_CASE_NAME, 'rate'),
+    [
+        ('0-dim', 0.4),
+    ],
+)
+class TestExponentialSampleKS(unittest.TestCase):
+    def setUp(self):
+        rate = paddle.to_tensor(self.rate, dtype=paddle.float32)
+        self.scale = rate.reciprocal()
+        self._paddle_expon = exponential.Exponential(rate)
+
+    def test_sample_ks(self):
+        sample_shape = (10000,)
+        samples = self._paddle_expon.sample(sample_shape)
+        self.assertTrue(self._kstest(samples))
+
+    def test_rsample_ks(self):
+        sample_shape = (10000,)
+        samples = self._paddle_expon.rsample(sample_shape)
+        self.assertTrue(self._kstest(samples))
+
+    def _kstest(self, samples):
+        # Uses the Kolmogorov-Smirnov test for goodness of fit.
+        ks, _ = scipy.stats.kstest(
+            samples, scipy.stats.expon(scale=self.scale).cdf
+        )
+        return ks < 0.02
+
+
+@parameterize.place(config.DEVICES)
+@parameterize.parameterize_cls(
+    (parameterize.TEST_CASE_NAME, 'rate1', 'rate2'),
+    [
+        (
+            'one-dim',
+            parameterize.xrand(
+                (2,),
+                dtype='float32',
+                min=np.finfo(dtype='float32').tiny,
+            ),
+            parameterize.xrand(
+                (2,),
+                dtype='float32',
+                min=np.finfo(dtype='float32').tiny,
+            ),
+        ),
+        (
+            'multi-dim',
+            parameterize.xrand(
+                (2, 3),
+                dtype='float32',
+                min=np.finfo(dtype='float32').tiny,
+            ),
+            parameterize.xrand(
+                (2, 3),
+                dtype='float32',
+                min=np.finfo(dtype='float32').tiny,
+            ),
+        ),
+    ],
+)
+class TestExponentialKL(unittest.TestCase):
+    def setUp(self):
+        self._expon1 = exponential.Exponential(paddle.to_tensor(self.rate1))
+        self._expon2 = exponential.Exponential(paddle.to_tensor(self.rate2))
+
+    def test_kl_divergence(self):
+        np.testing.assert_allclose(
+            kl.kl_divergence(self._expon1, self._expon2),
+            self._kl(),
+            rtol=config.RTOL.get(str(self._expon1.rate.numpy().dtype)),
+            atol=config.ATOL.get(str(self._expon1.rate.numpy().dtype)),
+        )
+
+    def test_kl1_error(self):
+        self.assertRaises(
+            TypeError,
+            self._expon1.kl_divergence,
+            paddle.distribution.beta.Beta,
+        )
+
+    def _kl(self):
+        rate_ratio = self.rate2 / self.rate1
+        t1 = -np.log(rate_ratio)
+        return t1 + rate_ratio - 1
+
+
+if __name__ == '__main__':
+    unittest.main()

test/distribution/test_distribution_exponential_static.py

@@ -0,0 +1,445 @@
+# 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
+#
+#     http://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
+import parameterize
+import scipy.stats
+from distribution import config
+
+import paddle
+from paddle.distribution import exponential
+
+np.random.seed(2023)
+paddle.seed(2023)
+
+paddle.enable_static()
+
+
+@parameterize.place(config.DEVICES)
+@parameterize.parameterize_cls(
+    (parameterize.TEST_CASE_NAME, 'rate'),
+    [
+        (
+            'one-dim',
+            parameterize.xrand(
+                (4,),
+                dtype='float32',
+                min=np.finfo(dtype='float32').tiny,
+            ),
+        ),
+        (
+            'multi-dim',
+            parameterize.xrand(
+                (10, 12),
+                dtype='float32',
+                min=np.finfo(dtype='float32').tiny,
+            ),
+        ),
+    ],
+)
+class TestExponential(unittest.TestCase):
+    def setUp(self):
+        self.program = paddle.static.Program()
+        self.executor = paddle.static.Executor(self.place)
+        with paddle.static.program_guard(self.program):
+            self.scale = 1 / self.rate
+            rate = paddle.static.data('rate', self.rate.shape, self.rate.dtype)
+            self._paddle_expon = exponential.Exponential(rate)
+            self.feeds = {'rate': self.rate}
+
+    def test_mean(self):
+        with paddle.static.program_guard(self.program):
+            [mean] = self.executor.run(
+                self.program,
+                feed=self.feeds,
+                fetch_list=[self._paddle_expon.mean],
+            )
+            np.testing.assert_allclose(
+                mean,
+                scipy.stats.expon.mean(scale=self.scale),
+                rtol=config.RTOL.get(str(self.rate.dtype)),
+                atol=config.ATOL.get(str(self.rate.dtype)),
+            )
+
+    def test_variance(self):
+        with paddle.static.program_guard(self.program):
+            [variance] = self.executor.run(
+                self.program,
+                feed=self.feeds,
+                fetch_list=[self._paddle_expon.variance],
+            )
+            np.testing.assert_allclose(
+                variance,
+                scipy.stats.expon.var(scale=self.scale),
+                rtol=config.RTOL.get(str(self.rate.dtype)),
+                atol=config.ATOL.get(str(self.rate.dtype)),
+            )
+
+    def test_entropy(self):
+        with paddle.static.program_guard(self.program):
+            [entropy] = self.executor.run(
+                self.program,
+                feed=self.feeds,
+                fetch_list=[self._paddle_expon.entropy()],
+            )
+            np.testing.assert_allclose(
+                entropy,
+                scipy.stats.expon.entropy(scale=self.scale),
+                rtol=config.RTOL.get(str(self.rate.dtype)),
+                atol=config.ATOL.get(str(self.rate.dtype)),
+            )
+
+    def test_prob(self):
+        with paddle.static.program_guard(self.program):
+            value = paddle.static.data(
+                'value',
+                self._paddle_expon.rate.shape,
+                self._paddle_expon.rate.dtype,
+            )
+            prob = self._paddle_expon.prob(value)
+
+            random_number = np.random.rand(
+                *self._paddle_expon.rate.shape
+            ).astype(self.rate.dtype)
+            feeds = dict(self.feeds, value=random_number)
+            [prob] = self.executor.run(
+                self.program, feed=feeds, fetch_list=[prob]
+            )
+            np.testing.assert_allclose(
+                prob,
+                scipy.stats.expon.pdf(random_number, scale=self.scale),
+                rtol=config.RTOL.get(str(self.rate.dtype)),
+                atol=config.ATOL.get(str(self.rate.dtype)),
+            )
+
+    def test_log_prob(self):
+        with paddle.static.program_guard(self.program):
+            value = paddle.static.data(
+                'value',
+                self._paddle_expon.rate.shape,
+                self._paddle_expon.rate.dtype,
+            )
+            log_prob = self._paddle_expon.log_prob(value)
+
+            random_number = np.random.rand(
+                *self._paddle_expon.rate.shape
+            ).astype(self.rate.dtype)
+            feeds = dict(self.feeds, value=random_number)
+            [log_prob] = self.executor.run(
+                self.program, feed=feeds, fetch_list=[log_prob]
+            )
+            np.testing.assert_allclose(
+                log_prob,
+                scipy.stats.expon.logpdf(random_number, scale=self.scale),
+                rtol=config.RTOL.get(str(self.rate.dtype)),
+                atol=config.ATOL.get(str(self.rate.dtype)),
+            )
+
+    def test_cdf(self):
+        with paddle.static.program_guard(self.program):
+            value = paddle.static.data(
+                'value',
+                self._paddle_expon.rate.shape,
+                self._paddle_expon.rate.dtype,
+            )
+            cdf = self._paddle_expon.cdf(value)
+
+            random_number = np.random.rand(
+                *self._paddle_expon.rate.shape
+            ).astype(self.rate.dtype)
+            feeds = dict(self.feeds, value=random_number)
+            [cdf] = self.executor.run(
+                self.program, feed=feeds, fetch_list=[cdf]
+            )
+            np.testing.assert_allclose(
+                cdf,
+                scipy.stats.expon.cdf(random_number, scale=self.scale),
+                rtol=config.RTOL.get(str(self.rate.dtype)),
+                atol=config.ATOL.get(str(self.rate.dtype)),
+            )
+
+    def test_icdf(self):
+        with paddle.static.program_guard(self.program):
+            value = paddle.static.data(
+                'value',
+                self._paddle_expon.rate.shape,
+                self._paddle_expon.rate.dtype,
+            )
+            icdf = self._paddle_expon.icdf(value)
+
+            random_number = np.random.rand(
+                *self._paddle_expon.rate.shape
+            ).astype(self.rate.dtype)
+            feeds = dict(self.feeds, value=random_number)
+            [icdf] = self.executor.run(
+                self.program, feed=feeds, fetch_list=[icdf]
+            )
+            np.testing.assert_allclose(
+                icdf,
+                scipy.stats.expon.ppf(random_number, scale=self.scale),
+                rtol=config.RTOL.get(str(self.rate.dtype)),
+                atol=config.ATOL.get(str(self.rate.dtype)),
+            )
+
+
+@parameterize.place(config.DEVICES)
+@parameterize.parameterize_cls(
+    (parameterize.TEST_CASE_NAME, 'rate'),
+    [
+        (
+            'one-dim',
+            parameterize.xrand(
+                (2,),
+                dtype='float32',
+                min=np.finfo(dtype='float32').tiny,
+            ),
+        ),
+        (
+            'multi-dim',
+            parameterize.xrand(
+                (2, 3),
+                dtype='float32',
+                min=np.finfo(dtype='float32').tiny,
+            ),
+        ),
+    ],
+)
+class TestExponentialSample(unittest.TestCase):
+    def setUp(self):
+        self.program = paddle.static.Program()
+        self.executor = paddle.static.Executor(self.place)
+        with paddle.static.program_guard(self.program):
+            self.scale = 1 / self.rate
+            rate = paddle.static.data('rate', self.rate.shape, self.rate.dtype)
+            self._paddle_expon = exponential.Exponential(rate)
+            self.feeds = {'rate': self.rate}
+
+    def test_sample_shape(self):
+        cases = [
+            {
+                'input': (),
+                'expect': tuple(np.squeeze(self.rate).shape),
+            },
+            {
+                'input': (4, 2),
+                'expect': (4, 2, *np.squeeze(self.rate).shape),
+            },
+        ]
+        for case in cases:
+            with paddle.static.program_guard(self.program):
+                [data] = self.executor.run(
+                    self.program,
+                    feed=self.feeds,
+                    fetch_list=self._paddle_expon.sample(case.get('input')),
+                )
+
+                self.assertTrue(data.shape == case.get('expect'))
+
+    def test_rsample_shape(self):
+        cases = [
+            {
+                'input': (),
+                'expect': tuple(np.squeeze(self.rate).shape),
+            },
+            {
+                'input': (3, 2),
+                'expect': (3, 2, *np.squeeze(self.rate).shape),
+            },
+        ]
+        for case in cases:
+            with paddle.static.program_guard(self.program):
+                [data] = self.executor.run(
+                    self.program,
+                    feed=self.feeds,
+                    fetch_list=self._paddle_expon.rsample(case.get('input')),
+                )
+
+                self.assertTrue(data.shape == case.get('expect'))
+
+    def test_sample(self):
+        sample_shape = (20000,)
+        with paddle.static.program_guard(self.program):
+            [data] = self.executor.run(
+                self.program,
+                feed=self.feeds,
+                fetch_list=self._paddle_expon.sample(sample_shape),
+            )
+            except_shape = sample_shape + np.squeeze(self.rate).shape
+            self.assertTrue(data.shape == except_shape)
+            np.testing.assert_allclose(
+                data.mean(axis=0),
+                scipy.stats.expon.mean(scale=self.scale),
+                rtol=0.1,
+                atol=config.ATOL.get(str(self.rate.dtype)),
+            )
+            np.testing.assert_allclose(
+                data.var(axis=0),
+                scipy.stats.expon.var(scale=self.scale),
+                rtol=0.1,
+                atol=config.ATOL.get(str(self.rate.dtype)),
+            )
+
+    def test_rsample(self):
+        sample_shape = (20000,)
+        with paddle.static.program_guard(self.program):
+            [data] = self.executor.run(
+                self.program,
+                feed=self.feeds,
+                fetch_list=self._paddle_expon.rsample(sample_shape),
+            )
+            except_shape = sample_shape + np.squeeze(self.rate).shape
+            self.assertTrue(data.shape == except_shape)
+            np.testing.assert_allclose(
+                data.mean(axis=0),
+                scipy.stats.expon.mean(scale=self.scale),
+                rtol=0.1,
+                atol=config.ATOL.get(str(self.rate.dtype)),
+            )
+            np.testing.assert_allclose(
+                data.var(axis=0),
+                scipy.stats.expon.var(scale=self.scale),
+                rtol=0.1,
+                atol=config.ATOL.get(str(self.rate.dtype)),
+            )
+
+
+@parameterize.place(config.DEVICES)
+@parameterize.parameterize_cls(
+    (parameterize.TEST_CASE_NAME, 'rate'),
+    [
+        ('0-dim', 0.4),
+    ],
+)
+class TestExponentialSampleKS(unittest.TestCase):
+    def setUp(self):
+        self.program = paddle.static.Program()
+        self.executor = paddle.static.Executor(self.place)
+        with paddle.static.program_guard(self.program):
+            self.scale = 1 / self.rate
+            rate = paddle.static.data('rate', (), 'float')
+            self._paddle_expon = exponential.Exponential(rate)
+            self.feeds = {'rate': self.rate}
+
+    def test_sample(self):
+        sample_shape = (10000,)
+        with paddle.static.program_guard(self.program):
+            [samples] = self.executor.run(
+                self.program,
+                feed=self.feeds,
+                fetch_list=self._paddle_expon.sample(sample_shape),
+            )
+            self.assertTrue(self._kstest(samples))
+
+    def test_rsample(self):
+        sample_shape = (10000,)
+        with paddle.static.program_guard(self.program):
+            [samples] = self.executor.run(
+                self.program,
+                feed=self.feeds,
+                fetch_list=self._paddle_expon.rsample(sample_shape),
+            )
+            self.assertTrue(self._kstest(samples))
+
+    def _kstest(self, samples):
+        # Uses the Kolmogorov-Smirnov test for goodness of fit.
+        ks, _ = scipy.stats.kstest(
+            samples, scipy.stats.expon(scale=self.scale).cdf
+        )
+        return ks < 0.02
+
+
+@parameterize.place(config.DEVICES)
+@parameterize.parameterize_cls(
+    (parameterize.TEST_CASE_NAME, 'rate1', 'rate2'),
+    [
+        (
+            'one-dim',
+            parameterize.xrand(
+                (2,),
+                dtype='float32',
+                min=np.finfo(dtype='float32').tiny,
+            ),
+            parameterize.xrand(
+                (2,),
+                dtype='float32',
+                min=np.finfo(dtype='float32').tiny,
+            ),
+        ),
+        (
+            'multi-dim',
+            parameterize.xrand(
+                (2, 3),
+                dtype='float32',
+                min=np.finfo(dtype='float32').tiny,
+            ),
+            parameterize.xrand(
+                (2, 3),
+                dtype='float32',
+                min=np.finfo(dtype='float32').tiny,
+            ),
+        ),
+    ],
+)
+class TestExponentialKL(unittest.TestCase):
+    def setUp(self):
+        self.program1 = paddle.static.Program()
+        self.program2 = paddle.static.Program()
+        self.executor = paddle.static.Executor(self.place)
+        with paddle.static.program_guard(self.program1, self.program2):
+            rate1 = paddle.static.data(
+                'rate1', self.rate1.shape, self.rate1.dtype
+            )
+            rate2 = paddle.static.data(
+                'rate2', self.rate2.shape, self.rate2.dtype
+            )
+
+            self._expon1 = exponential.Exponential(rate1)
+            self._expon2 = exponential.Exponential(rate2)
+
+            self.feeds = {
+                'rate1': self.rate1,
+                'rate2': self.rate2,
+            }
+
+    def test_kl_divergence(self):
+        with paddle.static.program_guard(self.program1, self.program2):
+            self.executor.run(self.program2)
+            [kl] = self.executor.run(
+                self.program1,
+                feed=self.feeds,
+                fetch_list=[self._expon1.kl_divergence(self._expon2)],
+            )
+            np.testing.assert_allclose(
+                kl,
+                self._kl(),
+                rtol=config.RTOL.get(str(self.rate1.dtype)),
+                atol=config.ATOL.get(str(self.rate1.dtype)),
+            )
+
+    def test_kl1_error(self):
+        self.assertRaises(
+            TypeError,
+            self._expon1.kl_divergence,
+            paddle.distribution.beta.Beta,
+        )
+
+    def _kl(self):
+        rate_ratio = self.rate2 / self.rate1
+        t1 = -np.log(rate_ratio)
+        return t1 + rate_ratio - 1
+
+
+if __name__ == '__main__':
+    unittest.main()

test/distribution/test_distribution_gamma.py

@@ -0,0 +1,467 @@
+# 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
+#
+#     http://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 numbers
+import unittest
+
+import numpy as np
+import parameterize
+import scipy.stats
+from distribution import config
+
+import paddle
+from paddle.distribution import gamma, kl
+
+np.random.seed(2023)
+paddle.seed(2023)
+
+
+@parameterize.place(config.DEVICES)
+@parameterize.parameterize_cls(
+    (parameterize.TEST_CASE_NAME, 'concentration', 'rate'),
+    [
+        (
+            '0-dim',
+            0.5,
+            0.5,
+        ),
+        (
+            'one-dim',
+            parameterize.xrand(
+                (6,),
+                dtype='float32',
+                min=np.finfo(dtype='float32').tiny,
+            ),
+            parameterize.xrand(
+                (6,),
+                dtype='float32',
+                min=np.finfo(dtype='float32').tiny,
+            ),
+        ),
+        (
+            'multi-dim',
+            parameterize.xrand(
+                (10, 12),
+                dtype='float32',
+                min=np.finfo(dtype='float32').tiny,
+            ),
+            parameterize.xrand(
+                (10, 12),
+                dtype='float32',
+                min=np.finfo(dtype='float32').tiny,
+            ),
+        ),
+        (
+            'broadcast',
+            parameterize.xrand(
+                (4, 1),
+                dtype='float32',
+                min=np.finfo(dtype='float32').tiny,
+            ),
+            parameterize.xrand(
+                (4, 6),
+                dtype='float32',
+                min=np.finfo(dtype='float32').tiny,
+            ),
+        ),
+    ],
+)
+class TestGamma(unittest.TestCase):
+    def setUp(self):
+        concentration = self.concentration
+        if not isinstance(self.concentration, numbers.Real):
+            concentration = paddle.to_tensor(self.concentration)
+
+        rate = self.rate
+        if not isinstance(self.rate, numbers.Real):
+            rate = paddle.to_tensor(self.rate)
+
+        self.scale = 1 / rate
+        self._paddle_gamma = gamma.Gamma(concentration, rate)
+
+    def test_mean(self):
+        with paddle.base.dygraph.guard(self.place):
+            np.testing.assert_allclose(
+                self._paddle_gamma.mean,
+                scipy.stats.gamma.mean(self.concentration, scale=self.scale),
+                rtol=config.RTOL.get(
+                    str(self._paddle_gamma.concentration.numpy().dtype)
+                ),
+                atol=config.ATOL.get(
+                    str(self._paddle_gamma.concentration.numpy().dtype)
+                ),
+            )
+
+    def test_variance(self):
+        with paddle.base.dygraph.guard(self.place):
+            np.testing.assert_allclose(
+                self._paddle_gamma.variance,
+                scipy.stats.gamma.var(self.concentration, scale=self.scale),
+                rtol=config.RTOL.get(
+                    str(self._paddle_gamma.concentration.numpy().dtype)
+                ),
+                atol=config.ATOL.get(
+                    str(self._paddle_gamma.concentration.numpy().dtype)
+                ),
+            )
+
+    def test_prob(self):
+        value = np.random.rand(*self._paddle_gamma.rate.shape)
+        with paddle.base.dygraph.guard(self.place):
+            np.testing.assert_allclose(
+                self._paddle_gamma.prob(paddle.to_tensor(value)),
+                scipy.stats.gamma.pdf(
+                    value, self.concentration, scale=self.scale
+                ),
+                rtol=config.RTOL.get(
+                    str(self._paddle_gamma.concentration.numpy().dtype)
+                ),
+                atol=config.ATOL.get(
+                    str(self._paddle_gamma.concentration.numpy().dtype)
+                ),
+            )
+
+    def test_log_prob(self):
+        value = np.random.rand(*self._paddle_gamma.rate.shape)
+        with paddle.base.dygraph.guard(self.place):
+            np.testing.assert_allclose(
+                self._paddle_gamma.log_prob(paddle.to_tensor(value)),
+                scipy.stats.gamma.logpdf(
+                    value, self.concentration, scale=self.scale
+                ),
+                rtol=config.RTOL.get(
+                    str(self._paddle_gamma.concentration.numpy().dtype)
+                ),
+                atol=config.ATOL.get(
+                    str(self._paddle_gamma.concentration.numpy().dtype)
+                ),
+            )
+
+    def test_entropy(self):
+        with paddle.base.dygraph.guard(self.place):
+            np.testing.assert_allclose(
+                self._paddle_gamma.entropy(),
+                scipy.stats.gamma.entropy(self.concentration, scale=self.scale),
+                rtol=config.RTOL.get(
+                    str(self._paddle_gamma.concentration.numpy().dtype)
+                ),
+                atol=config.ATOL.get(
+                    str(self._paddle_gamma.concentration.numpy().dtype)
+                ),
+            )
+
+
+@parameterize.place(config.DEVICES)
+@parameterize.parameterize_cls(
+    (parameterize.TEST_CASE_NAME, 'concentration', 'rate'),
+    [
+        (
+            'one-dim',
+            parameterize.xrand(
+                (2,),
+                dtype='float32',
+                min=np.finfo(dtype='float32').tiny,
+            ),
+            parameterize.xrand(
+                (2,),
+                dtype='float32',
+                min=np.finfo(dtype='float32').tiny,
+            ),
+        ),
+        (
+            'multi-dim',
+            parameterize.xrand(
+                (2, 3),
+                dtype='float32',
+                min=np.finfo(dtype='float32').tiny,
+            ),
+            parameterize.xrand(
+                (2, 3),
+                dtype='float32',
+                min=np.finfo(dtype='float32').tiny,
+            ),
+        ),
+    ],
+)
+class TestGammaSample(unittest.TestCase):
+    def setUp(self):
+        concentration = self.concentration
+        if not isinstance(self.concentration, numbers.Real):
+            concentration = paddle.to_tensor(self.concentration)
+
+        rate = self.rate
+        if not isinstance(self.rate, numbers.Real):
+            rate = paddle.to_tensor(self.rate)
+
+        self.scale = 1 / rate
+        self._paddle_gamma = gamma.Gamma(concentration, rate)
+
+    def test_sample_shape(self):
+        cases = [
+            {
+                'input': (),
+                'expect': tuple(paddle.squeeze(self._paddle_gamma.rate).shape),
+            },
+            {
+                'input': (3, 2),
+                'expect': (
+                    3,
+                    2,
+                    *paddle.squeeze(self._paddle_gamma.rate).shape,
+                ),
+            },
+        ]
+        for case in cases:
+            self.assertTrue(
+                tuple(self._paddle_gamma.sample(case.get('input')).shape)
+                == case.get('expect')
+            )
+
+    def test_rsample_shape(self):
+        cases = [
+            {
+                'input': (),
+                'expect': tuple(paddle.squeeze(self._paddle_gamma.rate).shape),
+            },
+            {
+                'input': (3, 2),
+                'expect': (
+                    3,
+                    2,
+                    *paddle.squeeze(self._paddle_gamma.rate).shape,
+                ),
+            },
+        ]
+        for case in cases:
+            self.assertTrue(
+                tuple(self._paddle_gamma.rsample(case.get('input')).shape)
+                == case.get('expect')
+            )
+
+    def test_sample(self):
+        sample_shape = (30000,)
+        samples = self._paddle_gamma.sample(sample_shape)
+        sample_values = samples.numpy()
+        self.assertEqual(sample_values.dtype, self.rate.dtype)
+
+        np.testing.assert_allclose(
+            sample_values.mean(axis=0),
+            scipy.stats.gamma.mean(self.concentration, scale=self.scale),
+            rtol=0.1,
+            atol=config.ATOL.get(
+                str(self._paddle_gamma.concentration.numpy().dtype)
+            ),
+        )
+        np.testing.assert_allclose(
+            sample_values.var(axis=0),
+            scipy.stats.gamma.var(self.concentration, scale=self.scale),
+            rtol=0.1,
+            atol=config.ATOL.get(
+                str(self._paddle_gamma.concentration.numpy().dtype)
+            ),
+        )
+
+    def test_rsample(self):
+        sample_shape = (30000,)
+        samples = self._paddle_gamma.rsample(sample_shape)
+        sample_values = samples.numpy()
+        self.assertEqual(sample_values.dtype, self.rate.dtype)
+
+        np.testing.assert_allclose(
+            sample_values.mean(axis=0),
+            scipy.stats.gamma.mean(self.concentration, scale=self.scale),
+            rtol=0.1,
+            atol=config.ATOL.get(
+                str(self._paddle_gamma.concentration.numpy().dtype)
+            ),
+        )
+        np.testing.assert_allclose(
+            sample_values.var(axis=0),
+            scipy.stats.gamma.var(self.concentration, scale=self.scale),
+            rtol=0.1,
+            atol=config.ATOL.get(
+                str(self._paddle_gamma.concentration.numpy().dtype)
+            ),
+        )
+
+    @unittest.skip("TODO: implement standard_gamma grad op.")
+    def test_rsample_backpropagation(self):
+        sample_shape = (1000,)
+        with paddle.base.dygraph.guard(self.place):
+            self._paddle_gamma.concentration.stop_gradient = False
+            self._paddle_gamma.rate.stop_gradient = False
+            samples = self._paddle_gamma.rsample(sample_shape)
+            grads = paddle.grad(
+                [samples],
+                [self._paddle_gamma.concentration, self._paddle_gamma.rate],
+            )
+            self.assertEqual(len(grads), 2)
+            self.assertEqual(
+                grads[0].dtype, self._paddle_gamma.concentration.dtype
+            )
+            self.assertEqual(
+                grads[0].shape, self._paddle_gamma.concentration.shape
+            )
+            self.assertEqual(grads[1].dtype, self._paddle_gamma.rate.dtype)
+            self.assertEqual(grads[1].shape, self._paddle_gamma.rate.shape)
+
+            samples.backward()
+            self.assertEqual(
+                list(self._paddle_gamma.concentration.gradient().shape),
+                self._paddle_gamma.concentration.shape,
+            )
+            self.assertEqual(
+                list(self._paddle_gamma.rate.gradient().shape),
+                self._paddle_gamma.rate.shape,
+            )
+
+
+@parameterize.place(config.DEVICES)
+@parameterize.parameterize_cls(
+    (parameterize.TEST_CASE_NAME, 'concentration', 'rate'),
+    [
+        ('0-dim', 0.4, 0.5),
+    ],
+)
+class TestGammaSampleKS(unittest.TestCase):
+    def setUp(self):
+        concentration = self.concentration
+        if not isinstance(self.concentration, numbers.Real):
+            concentration = paddle.to_tensor(self.concentration)
+
+        rate = self.rate
+        if not isinstance(self.rate, numbers.Real):
+            rate = paddle.to_tensor(self.rate)
+
+        self.scale = 1 / rate
+        self._paddle_gamma = gamma.Gamma(concentration, rate)
+
+    def test_sample_ks(self):
+        sample_shape = (15000,)
+        samples = self._paddle_gamma.sample(sample_shape)
+        self.assertTrue(self._kstest(samples))
+
+    def test_rsample_ks(self):
+        sample_shape = (15000,)
+        samples = self._paddle_gamma.rsample(sample_shape)
+        self.assertTrue(self._kstest(samples))
+
+    def _kstest(self, samples):
+        # Uses the Kolmogorov-Smirnov test for goodness of fit.
+        ks, _ = scipy.stats.kstest(
+            samples, scipy.stats.gamma(self.concentration, scale=self.scale).cdf
+        )
+        return ks < 0.02
+
+
+@parameterize.place(config.DEVICES)
+@parameterize.parameterize_cls(
+    (
+        parameterize.TEST_CASE_NAME,
+        'concentration1',
+        'rate1',
+        'concentration2',
+        'rate2',
+    ),
+    [
+        (
+            'one-dim',
+            parameterize.xrand(
+                (2,),
+                dtype='float32',
+                min=np.finfo(dtype='float32').tiny,
+            ),
+            parameterize.xrand(
+                (2,),
+                dtype='float32',
+                min=np.finfo(dtype='float32').tiny,
+            ),
+            parameterize.xrand(
+                (2,),
+                dtype='float32',
+                min=np.finfo(dtype='float32').tiny,
+            ),
+            parameterize.xrand(
+                (2,),
+                dtype='float32',
+                min=np.finfo(dtype='float32').tiny,
+            ),
+        ),
+        (
+            'multi-dim',
+            parameterize.xrand(
+                (2, 3),
+                dtype='float32',
+                min=np.finfo(dtype='float32').tiny,
+            ),
+            parameterize.xrand(
+                (2, 3),
+                dtype='float32',
+                min=np.finfo(dtype='float32').tiny,
+            ),
+            parameterize.xrand(
+                (2, 3),
+                dtype='float32',
+                min=np.finfo(dtype='float32').tiny,
+            ),
+            parameterize.xrand(
+                (2, 3),
+                dtype='float32',
+                min=np.finfo(dtype='float32').tiny,
+            ),
+        ),
+    ],
+)
+class TestGammaKL(unittest.TestCase):
+    def setUp(self):
+        self._gamma1 = gamma.Gamma(
+            paddle.to_tensor(self.concentration1), paddle.to_tensor(self.rate1)
+        )
+        self._gamma2 = gamma.Gamma(
+            paddle.to_tensor(self.concentration2), paddle.to_tensor(self.rate2)
+        )
+
+    def test_kl_divergence(self):
+        np.testing.assert_allclose(
+            kl.kl_divergence(self._gamma1, self._gamma2),
+            self._kl(),
+            rtol=config.RTOL.get(str(self._gamma1.concentration.numpy().dtype)),
+            atol=config.ATOL.get(str(self._gamma1.concentration.numpy().dtype)),
+        )
+
+    def test_kl1_error(self):
+        self.assertRaises(
+            TypeError,
+            self._gamma1.kl_divergence,
+            paddle.distribution.beta.Beta,
+        )
+
+    def _kl(self):
+        concentration1 = self.concentration1
+        concentration2 = self.concentration2
+        rate1 = self.rate1
+        rate2 = self.rate2
+        t1 = concentration2 * np.log(rate1 / rate2)
+        t2 = scipy.special.gammaln(concentration2) - scipy.special.gammaln(
+            concentration1
+        )
+        t3 = (concentration1 - concentration2) * scipy.special.digamma(
+            concentration1
+        )
+        t4 = (rate2 - rate1) * (concentration1 / rate1)
+        return t1 + t2 + t3 + t4
+
+
+if __name__ == '__main__':
+    unittest.main()

test/distribution/test_distribution_gamma_static.py

@@ -0,0 +1,509 @@
+# 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
+#
+#     http://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
+import parameterize
+import scipy.stats
+from distribution import config
+
+import paddle
+from paddle.distribution import gamma
+
+np.random.seed(2023)
+paddle.seed(2023)
+
+paddle.enable_static()
+
+
+@parameterize.place(config.DEVICES)
+@parameterize.parameterize_cls(
+    (parameterize.TEST_CASE_NAME, 'concentration', 'rate'),
+    [
+        (
+            'one-dim',
+            parameterize.xrand(
+                (6,),
+                dtype='float32',
+                min=np.finfo(dtype='float32').tiny,
+            ),
+            parameterize.xrand(
+                (6,),
+                dtype='float32',
+                min=np.finfo(dtype='float32').tiny,
+            ),
+        ),
+        (
+            'multi-dim',
+            parameterize.xrand(
+                (10, 12),
+                dtype='float32',
+                min=np.finfo(dtype='float32').tiny,
+            ),
+            parameterize.xrand(
+                (10, 12),
+                dtype='float32',
+                min=np.finfo(dtype='float32').tiny,
+            ),
+        ),
+        (
+            'broadcast',
+            parameterize.xrand(
+                (4, 1),
+                dtype='float32',
+                min=np.finfo(dtype='float32').tiny,
+            ),
+            parameterize.xrand(
+                (4, 6),
+                dtype='float32',
+                min=np.finfo(dtype='float32').tiny,
+            ),
+        ),
+    ],
+)
+class TestGamma(unittest.TestCase):
+    def setUp(self):
+        self.program = paddle.static.Program()
+        self.executor = paddle.static.Executor(self.place)
+        with paddle.static.program_guard(self.program):
+            self.scale = 1 / self.rate
+            concentration = paddle.static.data(
+                'concentration',
+                self.concentration.shape,
+                self.concentration.dtype,
+            )
+            rate = paddle.static.data('rate', self.rate.shape, self.rate.dtype)
+            self._paddle_gamma = gamma.Gamma(concentration, rate)
+            self.feeds = {
+                'concentration': self.concentration,
+                'rate': self.rate,
+            }
+
+    def test_mean(self):
+        with paddle.static.program_guard(self.program):
+            [mean] = self.executor.run(
+                self.program,
+                feed=self.feeds,
+                fetch_list=[self._paddle_gamma.mean],
+            )
+            np.testing.assert_allclose(
+                mean,
+                scipy.stats.gamma.mean(self.concentration, scale=self.scale),
+                rtol=config.RTOL.get(str(self.concentration.dtype)),
+                atol=config.ATOL.get(str(self.concentration.dtype)),
+            )
+
+    def test_variance(self):
+        with paddle.static.program_guard(self.program):
+            [variance] = self.executor.run(
+                self.program,
+                feed=self.feeds,
+                fetch_list=[self._paddle_gamma.variance],
+            )
+            np.testing.assert_allclose(
+                variance,
+                scipy.stats.gamma.var(self.concentration, scale=self.scale),
+                rtol=config.RTOL.get(str(self.concentration.dtype)),
+                atol=config.ATOL.get(str(self.concentration.dtype)),
+            )
+
+    def test_entropy(self):
+        with paddle.static.program_guard(self.program):
+            [entropy] = self.executor.run(
+                self.program,
+                feed=self.feeds,
+                fetch_list=[self._paddle_gamma.entropy()],
+            )
+            np.testing.assert_allclose(
+                entropy,
+                scipy.stats.gamma.entropy(self.concentration, scale=self.scale),
+                rtol=config.RTOL.get(str(self.concentration.dtype)),
+                atol=config.ATOL.get(str(self.concentration.dtype)),
+            )
+
+    def test_prob(self):
+        with paddle.static.program_guard(self.program):
+            value = paddle.static.data(
+                'value',
+                self._paddle_gamma.concentration.shape,
+                self._paddle_gamma.concentration.dtype,
+            )
+            prob = self._paddle_gamma.prob(value)
+
+            random_number = np.random.rand(
+                *self._paddle_gamma.concentration.shape
+            ).astype(self.concentration.dtype)
+            feeds = dict(self.feeds, value=random_number)
+            [prob] = self.executor.run(
+                self.program, feed=feeds, fetch_list=[prob]
+            )
+            np.testing.assert_allclose(
+                prob,
+                scipy.stats.gamma.pdf(
+                    random_number, self.concentration, scale=self.scale
+                ),
+                rtol=config.RTOL.get(str(self.concentration.dtype)),
+                atol=config.ATOL.get(str(self.concentration.dtype)),
+            )
+
+    def test_log_prob(self):
+        with paddle.static.program_guard(self.program):
+            value = paddle.static.data(
+                'value',
+                self._paddle_gamma.concentration.shape,
+                self._paddle_gamma.concentration.dtype,
+            )
+            log_prob = self._paddle_gamma.log_prob(value)
+
+            random_number = np.random.rand(
+                *self._paddle_gamma.concentration.shape
+            ).astype(self.concentration.dtype)
+            feeds = dict(self.feeds, value=random_number)
+            [log_prob] = self.executor.run(
+                self.program, feed=feeds, fetch_list=[log_prob]
+            )
+            np.testing.assert_allclose(
+                log_prob,
+                scipy.stats.gamma.logpdf(
+                    random_number, self.concentration, scale=self.scale
+                ),
+                rtol=config.RTOL.get(str(self.concentration.dtype)),
+                atol=config.ATOL.get(str(self.concentration.dtype)),
+            )
+
+
+@parameterize.place(config.DEVICES)
+@parameterize.parameterize_cls(
+    (parameterize.TEST_CASE_NAME, 'concentration', 'rate'),
+    [
+        (
+            'one-dim',
+            parameterize.xrand(
+                (2,),
+                dtype='float32',
+                min=np.finfo(dtype='float32').tiny,
+            ),
+            parameterize.xrand(
+                (2,),
+                dtype='float32',
+                min=np.finfo(dtype='float32').tiny,
+            ),
+        ),
+        (
+            'multi-dim',
+            parameterize.xrand(
+                (2, 3),
+                dtype='float32',
+                min=np.finfo(dtype='float32').tiny,
+            ),
+            parameterize.xrand(
+                (2, 3),
+                dtype='float32',
+                min=np.finfo(dtype='float32').tiny,
+            ),
+        ),
+    ],
+)
+class TestGammaSample(unittest.TestCase):
+    def setUp(self):
+        self.program = paddle.static.Program()
+        self.executor = paddle.static.Executor(self.place)
+        with paddle.static.program_guard(self.program):
+            self.scale = 1 / self.rate
+            concentration = paddle.static.data(
+                'concentration',
+                self.concentration.shape,
+                self.concentration.dtype,
+            )
+            rate = paddle.static.data('rate', self.rate.shape, self.rate.dtype)
+            self._paddle_gamma = gamma.Gamma(concentration, rate)
+            self.feeds = {
+                'concentration': self.concentration,
+                'rate': self.rate,
+            }
+
+    def test_sample_shape(self):
+        cases = [
+            {
+                'input': (),
+                'expect': tuple(np.squeeze(self.rate).shape),
+            },
+            {
+                'input': (4, 2),
+                'expect': (4, 2, *np.squeeze(self.rate).shape),
+            },
+        ]
+        for case in cases:
+            with paddle.static.program_guard(self.program):
+                [data] = self.executor.run(
+                    self.program,
+                    feed=self.feeds,
+                    fetch_list=self._paddle_gamma.sample(case.get('input')),
+                )
+
+                self.assertTrue(data.shape == case.get('expect'))
+
+    def test_rsample_shape(self):
+        cases = [
+            {
+                'input': (),
+                'expect': tuple(np.squeeze(self.rate).shape),
+            },
+            {
+                'input': (3, 2),
+                'expect': (3, 2, *np.squeeze(self.rate).shape),
+            },
+        ]
+        for case in cases:
+            with paddle.static.program_guard(self.program):
+                [data] = self.executor.run(
+                    self.program,
+                    feed=self.feeds,
+                    fetch_list=self._paddle_gamma.rsample(case.get('input')),
+                )
+
+                self.assertTrue(data.shape == case.get('expect'))
+
+    def test_sample(self):
+        sample_shape = (30000,)
+        with paddle.static.program_guard(self.program):
+            [data] = self.executor.run(
+                self.program,
+                feed=self.feeds,
+                fetch_list=self._paddle_gamma.sample(sample_shape),
+            )
+            except_shape = sample_shape + np.squeeze(self.rate).shape
+            self.assertTrue(data.shape == except_shape)
+            np.testing.assert_allclose(
+                data.mean(axis=0),
+                scipy.stats.gamma.mean(self.concentration, scale=self.scale),
+                rtol=0.1,
+                atol=config.ATOL.get(str(self.concentration.dtype)),
+            )
+            np.testing.assert_allclose(
+                data.var(axis=0),
+                scipy.stats.gamma.var(self.concentration, scale=self.scale),
+                rtol=0.1,
+                atol=config.ATOL.get(str(self.concentration.dtype)),
+            )
+
+    def test_rsample(self):
+        sample_shape = (30000,)
+        with paddle.static.program_guard(self.program):
+            [data] = self.executor.run(
+                self.program,
+                feed=self.feeds,
+                fetch_list=self._paddle_gamma.rsample(sample_shape),
+            )
+            except_shape = sample_shape + np.squeeze(self.rate).shape
+            self.assertTrue(data.shape == except_shape)
+            np.testing.assert_allclose(
+                data.mean(axis=0),
+                scipy.stats.gamma.mean(self.concentration, scale=self.scale),
+                rtol=0.1,
+                atol=config.ATOL.get(str(self.concentration.dtype)),
+            )
+            np.testing.assert_allclose(
+                data.var(axis=0),
+                scipy.stats.gamma.var(self.concentration, scale=self.scale),
+                rtol=0.1,
+                atol=config.ATOL.get(str(self.concentration.dtype)),
+            )
+
+
+@parameterize.place(config.DEVICES)
+@parameterize.parameterize_cls(
+    (parameterize.TEST_CASE_NAME, 'concentration', 'rate'),
+    [
+        ('0-dim', 0.4, 0.5),
+    ],
+)
+class TestGammaSampleKS(unittest.TestCase):
+    def setUp(self):
+        self.program = paddle.static.Program()
+        self.executor = paddle.static.Executor(self.place)
+        with paddle.static.program_guard(self.program):
+            self.scale = 1 / self.rate
+            concentration = paddle.static.data(
+                'concentration',
+                (),
+                'float',
+            )
+            rate = paddle.static.data('rate', (), 'float')
+            self._paddle_gamma = gamma.Gamma(concentration, rate)
+            self.feeds = {
+                'concentration': self.concentration,
+                'rate': self.rate,
+            }
+
+    def test_sample(self):
+        sample_shape = (15000,)
+        with paddle.static.program_guard(self.program):
+            [samples] = self.executor.run(
+                self.program,
+                feed=self.feeds,
+                fetch_list=self._paddle_gamma.sample(sample_shape),
+            )
+            self.assertTrue(self._kstest(samples))
+
+    def test_rsample(self):
+        sample_shape = (15000,)
+        with paddle.static.program_guard(self.program):
+            [samples] = self.executor.run(
+                self.program,
+                feed=self.feeds,
+                fetch_list=self._paddle_gamma.rsample(sample_shape),
+            )
+            self.assertTrue(self._kstest(samples))
+
+    def _kstest(self, samples):
+        # Uses the Kolmogorov-Smirnov test for goodness of fit.
+        ks, _ = scipy.stats.kstest(
+            samples, scipy.stats.gamma(self.concentration, scale=self.scale).cdf
+        )
+        return ks < 0.02
+
+
+@parameterize.place(config.DEVICES)
+@parameterize.parameterize_cls(
+    (
+        parameterize.TEST_CASE_NAME,
+        'concentration1',
+        'rate1',
+        'concentration2',
+        'rate2',
+    ),
+    [
+        (
+            'one-dim',
+            parameterize.xrand(
+                (2,),
+                dtype='float32',
+                min=np.finfo(dtype='float32').tiny,
+            ),
+            parameterize.xrand(
+                (2,),
+                dtype='float32',
+                min=np.finfo(dtype='float32').tiny,
+            ),
+            parameterize.xrand(
+                (2,),
+                dtype='float32',
+                min=np.finfo(dtype='float32').tiny,
+            ),
+            parameterize.xrand(
+                (2,),
+                dtype='float32',
+                min=np.finfo(dtype='float32').tiny,
+            ),
+        ),
+        (
+            'multi-dim',
+            parameterize.xrand(
+                (2, 3),
+                dtype='float32',
+                min=np.finfo(dtype='float32').tiny,
+            ),
+            parameterize.xrand(
+                (2, 3),
+                dtype='float32',
+                min=np.finfo(dtype='float32').tiny,
+            ),
+            parameterize.xrand(
+                (2, 3),
+                dtype='float32',
+                min=np.finfo(dtype='float32').tiny,
+            ),
+            parameterize.xrand(
+                (2, 3),
+                dtype='float32',
+                min=np.finfo(dtype='float32').tiny,
+            ),
+        ),
+    ],
+)
+class TestGammaKL(unittest.TestCase):
+    def setUp(self):
+        self.program1 = paddle.static.Program()
+        self.program2 = paddle.static.Program()
+        self.executor = paddle.static.Executor(self.place)
+        with paddle.static.program_guard(self.program1, self.program2):
+            concentration1 = paddle.static.data(
+                'concentration1',
+                self.concentration1.shape,
+                self.concentration1.dtype,
+            )
+            concentration2 = paddle.static.data(
+                'concentration2',
+                self.concentration2.shape,
+                self.concentration2.dtype,
+            )
+            rate1 = paddle.static.data(
+                'rate1', self.rate1.shape, self.rate1.dtype
+            )
+            rate2 = paddle.static.data(
+                'rate2', self.rate2.shape, self.rate2.dtype
+            )
+
+            self._gamma1 = gamma.Gamma(concentration1, rate1)
+            self._gamma2 = gamma.Gamma(concentration2, rate2)
+
+            self.feeds = {
+                'concentration1': self.concentration1,
+                'concentration2': self.concentration2,
+                'rate1': self.rate1,
+                'rate2': self.rate2,
+            }
+
+    def test_kl_divergence(self):
+        with paddle.static.program_guard(self.program1, self.program2):
+            self.executor.run(self.program2)
+            [kl] = self.executor.run(
+                self.program1,
+                feed=self.feeds,
+                fetch_list=[self._gamma1.kl_divergence(self._gamma2)],
+            )
+            np.testing.assert_allclose(
+                kl,
+                self._kl(),
+                rtol=config.RTOL.get(str(self.concentration1.dtype)),
+                atol=config.ATOL.get(str(self.concentration1.dtype)),
+            )
+
+    def test_kl1_error(self):
+        self.assertRaises(
+            TypeError,
+            self._gamma1.kl_divergence,
+            paddle.distribution.beta.Beta,
+        )
+
+    def _kl(self):
+        concentration1 = self.concentration1
+        concentration2 = self.concentration2
+        rate1 = self.rate1
+        rate2 = self.rate2
+        t1 = concentration2 * np.log(rate1 / rate2)
+        t2 = scipy.special.gammaln(concentration2) - scipy.special.gammaln(
+            concentration1
+        )
+        t3 = (concentration1 - concentration2) * scipy.special.digamma(
+            concentration1
+        )
+        t4 = (rate2 - rate1) * (concentration1 / rate1)
+        return t1 + t2 + t3 + t4
+
+
+if __name__ == '__main__':
+    unittest.main()

test/distribution/test_distribution_geometric.py

@@ -0,0 +1,350 @@
+# 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
+#
+#     http://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 numbers
+import unittest
+
+import numpy as np
+import scipy.stats
+from distribution.config import ATOL, DEVICES, RTOL
+from parameterize import TEST_CASE_NAME, parameterize_cls, place, xrand
+
+import paddle
+from paddle.distribution import geometric, kl
+from paddle.nn.functional import log_softmax
+
+np.random.seed(2023)
+
+
+@place(DEVICES)
+@parameterize_cls(
+    (TEST_CASE_NAME, 'probs'),
+    [
+        (
+            'one-dim',
+            xrand(
+                (2,),
+                dtype='float32',
+                min=np.finfo(dtype='float32').tiny,
+                max=1.0,
+            ),
+        ),
+        (
+            'multi-dim',
+            xrand(
+                (2, 3),
+                dtype='float32',
+                min=np.finfo(dtype='float32').tiny,
+                max=1.0,
+            ),
+        ),
+    ],
+)
+class TestGeometric(unittest.TestCase):
+    def setUp(self):
+        probs = self.probs
+        if not isinstance(self.probs, numbers.Real):
+            probs = paddle.to_tensor(self.probs, dtype=paddle.float32)
+
+        self._paddle_geom = geometric.Geometric(probs)
+
+    def test_mean(self):
+        with paddle.base.dygraph.guard(self.place):
+            np.testing.assert_allclose(
+                self._paddle_geom.mean,
+                scipy.stats.geom.mean(self.probs, loc=-1),
+                rtol=RTOL.get(str(self._paddle_geom.probs.numpy().dtype)),
+                atol=ATOL.get(str(self._paddle_geom.probs.numpy().dtype)),
+            )
+
+    def test_variance(self):
+        with paddle.base.dygraph.guard(self.place):
+            np.testing.assert_allclose(
+                self._paddle_geom.variance,
+                scipy.stats.geom.var(self.probs, loc=-1),
+                rtol=RTOL.get(str(self._paddle_geom.probs.numpy().dtype)),
+                atol=ATOL.get(str(self._paddle_geom.probs.numpy().dtype)),
+            )
+
+    def test_stddev(self):
+        with paddle.base.dygraph.guard(self.place):
+            np.testing.assert_allclose(
+                self._paddle_geom.stddev,
+                scipy.stats.geom.std(self.probs, loc=-1),
+                rtol=RTOL.get(str(self._paddle_geom.probs.numpy().dtype)),
+                atol=ATOL.get(str(self._paddle_geom.probs.numpy().dtype)),
+            )
+
+    def test_entropy(self):
+        with paddle.base.dygraph.guard(self.place):
+            np.testing.assert_allclose(
+                self._paddle_geom.entropy(),
+                scipy.stats.geom.entropy(self.probs, loc=-1),
+                rtol=RTOL.get(str(self._paddle_geom.probs.numpy().dtype)),
+                atol=ATOL.get(str(self._paddle_geom.probs.numpy().dtype)),
+            )
+
+    def test_init_prob_type_error(self):
+        with self.assertRaises(TypeError):
+            paddle.distribution.geometric.Geometric([2])
+
+    def test_sample_shape(self):
+        cases = [
+            {
+                'input': (),
+                'expect': tuple(paddle.squeeze(self._paddle_geom.probs).shape),
+            },
+            {
+                'input': (4, 2),
+                'expect': (
+                    4,
+                    2,
+                    *paddle.squeeze(self._paddle_geom.probs).shape,
+                ),
+            },
+        ]
+        for case in cases:
+            self.assertTrue(
+                tuple(self._paddle_geom.sample(case.get('input')).shape)
+                == case.get('expect')
+            )
+
+    def test_sample(self):
+        sample_shape = (100000,)
+        samples = self._paddle_geom.sample(sample_shape)
+        sample_values = samples.numpy()
+        self.assertEqual(sample_values.dtype, self.probs.dtype)
+
+        np.testing.assert_allclose(
+            sample_values.mean(axis=0),
+            scipy.stats.geom.mean(self.probs, loc=-1),
+            rtol=0.1,
+            atol=ATOL.get(str(self._paddle_geom.probs.numpy().dtype)),
+        )
+        np.testing.assert_allclose(
+            sample_values.var(axis=0),
+            scipy.stats.geom.var(self.probs, loc=-1),
+            rtol=0.1,
+            atol=ATOL.get(str(self._paddle_geom.probs.numpy().dtype)),
+        )
+
+    def test_rsample_shape(self):
+        cases = [
+            {
+                'input': (),
+                'expect': tuple(paddle.squeeze(self._paddle_geom.probs).shape),
+            },
+            {
+                'input': (2, 5),
+                'expect': (
+                    2,
+                    5,
+                    *paddle.squeeze(self._paddle_geom.probs).shape,
+                ),
+            },
+        ]
+        for case in cases:
+            self.assertTrue(
+                tuple(self._paddle_geom.rsample(case.get('input')).shape)
+                == case.get('expect')
+            )
+
+    def test_rsample(self):
+        sample_shape = (100000,)
+        samples = self._paddle_geom.rsample(sample_shape)
+        sample_values = samples.numpy()
+        self.assertEqual(sample_values.dtype, self.probs.dtype)
+
+        np.testing.assert_allclose(
+            sample_values.mean(axis=0),
+            scipy.stats.geom.mean(self.probs, loc=-1),
+            rtol=0.1,
+            atol=ATOL.get(str(self._paddle_geom.probs.numpy().dtype)),
+        )
+        np.testing.assert_allclose(
+            sample_values.var(axis=0),
+            scipy.stats.geom.var(self.probs, loc=-1),
+            rtol=0.1,
+            atol=ATOL.get(str(self._paddle_geom.probs.numpy().dtype)),
+        )
+
+    def test_back_rsample(self):
+        sample_shape = (100000,)
+
+        with paddle.base.dygraph.guard(self.place):
+            self._paddle_geom.probs.stop_gradient = False
+
+            rs_value = self._paddle_geom.rsample(sample_shape)
+            softmax_rs = log_softmax(rs_value)
+            grads = paddle.grad([softmax_rs], [self._paddle_geom.probs])
+
+            self.assertEqual(len(grads), 1)
+            self.assertEqual(grads[0].dtype, self._paddle_geom.probs.dtype)
+            self.assertEqual(grads[0].shape, self._paddle_geom.probs.shape)
+
+
+@place(DEVICES)
+@parameterize_cls(
+    (TEST_CASE_NAME, 'probs', 'value'),
+    [
+        (
+            'one-dim',
+            xrand(
+                (2,),
+                dtype='float32',
+                min=np.finfo(dtype='float32').tiny,
+                max=1.0,
+            ),
+            5,
+        ),
+        (
+            'mult-dim',
+            xrand(
+                (2, 2),
+                dtype='float32',
+                min=np.finfo(dtype='float32').tiny,
+                max=1.0,
+            ),
+            5,
+        ),
+        (
+            'mult-dim',
+            xrand(
+                (2, 2, 2),
+                dtype='float32',
+                min=np.finfo(dtype='float32').tiny,
+                max=1.0,
+            ),
+            5,
+        ),
+    ],
+)
+class TestGeometricPMF(unittest.TestCase):
+    def setUp(self):
+        self._paddle_geom = geometric.Geometric(
+            probs=paddle.to_tensor(self.probs)
+        )
+
+    def test_pmf(self):
+        with paddle.base.dygraph.guard(self.place):
+            np.testing.assert_allclose(
+                self._paddle_geom.pmf(self.value),
+                scipy.stats.geom.pmf(self.value, self.probs, loc=-1),
+                rtol=RTOL.get(str(self.probs.dtype)),
+                atol=ATOL.get(str(self.probs.dtype)),
+            )
+
+    def test_log_pmf(self):
+        with paddle.base.dygraph.guard(self.place):
+            np.testing.assert_allclose(
+                self._paddle_geom.log_pmf(self.value),
+                scipy.stats.geom.logpmf(self.value, self.probs, loc=-1),
+                rtol=RTOL.get(str(self.probs.dtype)),
+                atol=ATOL.get(str(self.probs.dtype)),
+            )
+
+    def test_cdf(self):
+        with paddle.base.dygraph.guard(self.place):
+            np.testing.assert_allclose(
+                self._paddle_geom.cdf(self.value),
+                scipy.stats.geom.cdf(self.value, self.probs, loc=-1),
+                rtol=RTOL.get(str(self._paddle_geom.probs.numpy().dtype)),
+                atol=ATOL.get(str(self._paddle_geom.probs.numpy().dtype)),
+            )
+
+    def test_pmf_error(self):
+        self.assertRaises(TypeError, self._paddle_geom.pmf, [1, 2])
+
+    def test_log_pmf_error(self):
+        self.assertRaises(TypeError, self._paddle_geom.log_pmf, [1, 2])
+
+    def test_cdf_error(self):
+        self.assertRaises(TypeError, self._paddle_geom.cdf, [1, 2])
+
+
+@place(DEVICES)
+@parameterize_cls(
+    (TEST_CASE_NAME, 'probs1', 'probs2'),
+    [
+        (
+            'one-dim',
+            xrand(
+                (2,),
+                dtype='float32',
+                min=np.finfo(dtype='float32').tiny,
+                max=1.0,
+            ),
+            xrand(
+                (2,),
+                dtype='float32',
+                min=np.finfo(dtype='float32').tiny,
+                max=1.0,
+            ),
+        ),
+        (
+            'multi-dim',
+            xrand(
+                (2, 2),
+                dtype='float32',
+                min=np.finfo(dtype='float32').tiny,
+                max=1.0,
+            ),
+            xrand(
+                (2, 2),
+                dtype='float32',
+                min=np.finfo(dtype='float32').tiny,
+                max=1.0,
+            ),
+        ),
+    ],
+)
+class TestGeometricKL(unittest.TestCase):
+    def setUp(self):
+        paddle.disable_static()
+        self._geometric1 = geometric.Geometric(
+            probs=paddle.to_tensor(self.probs1)
+        )
+        self._geometric2 = geometric.Geometric(
+            probs=paddle.to_tensor(self.probs2)
+        )
+
+    def test_kl_divergence(self):
+        np.testing.assert_allclose(
+            kl.kl_divergence(self._geometric1, self._geometric2),
+            self._kl(),
+            rtol=RTOL.get(str(self._geometric1.probs.numpy().dtype)),
+            atol=ATOL.get(str(self._geometric1.probs.numpy().dtype)),
+        )
+
+    def test_kl1_error(self):
+        self.assertRaises(
+            TypeError,
+            self._geometric1.kl_divergence,
+            paddle.distribution.beta.Beta,
+        )
+
+    def test_kl2_error(self):
+        self.assertRaises(
+            TypeError,
+            self._geometric2.kl_divergence,
+            paddle.distribution.beta.Beta,
+        )
+
+    def _kl(self):
+        return self.probs1 * np.log(self.probs1 / self.probs2) + (
+            1.0 - self.probs1
+        ) * np.log((1.0 - self.probs1) / (1.0 - self.probs2))
+
+
+if __name__ == '__main__':
+    unittest.main()

test/distribution/test_distribution_geometric_static.py

@@ -0,0 +1,346 @@
+# 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
+#
+#     http://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
+import scipy.stats
+from distribution.config import ATOL, DEVICES, RTOL
+from parameterize import TEST_CASE_NAME, parameterize_cls, place, xrand
+
+import paddle
+from paddle.distribution import geometric
+
+np.random.seed(2023)
+
+paddle.enable_static()
+
+
+@place(DEVICES)
+@parameterize_cls(
+    (TEST_CASE_NAME, 'probs'),
+    [
+        (
+            'one-dim',
+            xrand(
+                (2,),
+                dtype='float32',
+                min=np.finfo(dtype='float32').tiny,
+                max=1.0,
+            ),
+        ),
+        (
+            'multi-dim',
+            xrand(
+                (2, 3),
+                dtype='float32',
+                min=np.finfo(dtype='float32').tiny,
+                max=1.0,
+            ),
+        ),
+    ],
+)
+class TestGeometric(unittest.TestCase):
+    def setUp(self):
+        self.program = paddle.static.Program()
+        self.executor = paddle.static.Executor(self.place)
+        with paddle.static.program_guard(self.program):
+            # scale no need convert to tensor for scale input unittest
+            probs = paddle.static.data(
+                'probs', self.probs.shape, self.probs.dtype
+            )
+
+            self._paddle_geometric = geometric.Geometric(probs)
+            self.feeds = {'probs': self.probs}
+
+    def test_mean(self):
+        with paddle.static.program_guard(self.program):
+            [mean] = self.executor.run(
+                self.program,
+                feed=self.feeds,
+                fetch_list=[self._paddle_geometric.mean],
+            )
+            np.testing.assert_allclose(
+                mean,
+                scipy.stats.geom.mean(self.probs, loc=-1),
+                rtol=RTOL.get(str(self.probs.dtype)),
+                atol=ATOL.get(str(self.probs.dtype)),
+            )
+
+    def test_variance(self):
+        with paddle.static.program_guard(self.program):
+            [variance] = self.executor.run(
+                self.program,
+                feed=self.feeds,
+                fetch_list=[self._paddle_geometric.variance],
+            )
+            np.testing.assert_allclose(
+                variance,
+                scipy.stats.geom.var(self.probs, loc=-1),
+                rtol=RTOL.get(str(self.probs.dtype)),
+                atol=ATOL.get(str(self.probs.dtype)),
+            )
+
+    def test_stddev(self):
+        with paddle.static.program_guard(self.program):
+            [stddev] = self.executor.run(
+                self.program,
+                feed=self.feeds,
+                fetch_list=[self._paddle_geometric.stddev],
+            )
+            np.testing.assert_allclose(
+                stddev,
+                scipy.stats.geom.std(self.probs, loc=-1),
+                rtol=RTOL.get(str(self.probs.dtype)),
+                atol=ATOL.get(str(self.probs.dtype)),
+            )
+
+    def test_sample(self):
+        with paddle.static.program_guard(self.program):
+            [data] = self.executor.run(
+                self.program,
+                feed=self.feeds,
+                fetch_list=self._paddle_geometric.sample(),
+            )
+            self.assertTrue(
+                data.shape == np.broadcast_arrays(self.probs)[0].shape
+            )
+
+    def test_rsample(self):
+        with paddle.static.program_guard(self.program):
+            [data] = self.executor.run(
+                self.program,
+                feed=self.feeds,
+                fetch_list=self._paddle_geometric.rsample(),
+            )
+            self.assertTrue(
+                data.shape == np.broadcast_arrays(self.probs)[0].shape
+            )
+
+    def test_entropy(self):
+        with paddle.static.program_guard(self.program):
+            [entropy] = self.executor.run(
+                self.program,
+                feed=self.feeds,
+                fetch_list=[self._paddle_geometric.entropy()],
+            )
+            np.testing.assert_allclose(
+                entropy,
+                scipy.stats.geom.entropy(self.probs, loc=-1),
+                rtol=RTOL.get(str(self.probs.dtype)),
+                atol=ATOL.get(str(self.probs.dtype)),
+            )
+
+    def test_init_prob_type_error(self):
+        with self.assertRaises(TypeError):
+            paddle.distribution.geometric.Geometric([0.5])
+
+
+@place(DEVICES)
+@parameterize_cls(
+    (TEST_CASE_NAME, 'probs', 'value'),
+    [
+        (
+            'one-dim',
+            xrand(
+                (2,),
+                dtype='float32',
+                min=np.finfo(dtype='float32').tiny,
+                max=1.0,
+            ),
+            5,
+        ),
+        (
+            'mult-dim',
+            xrand(
+                (2, 2),
+                dtype='float32',
+                min=np.finfo(dtype='float32').tiny,
+                max=1.0,
+            ),
+            5,
+        ),
+        (
+            'mult-dim',
+            xrand(
+                (2, 2, 2),
+                dtype='float32',
+                min=np.finfo(dtype='float32').tiny,
+                max=1.0,
+            ),
+            5,
+        ),
+    ],
+)
+class TestGeometricPMF(unittest.TestCase):
+    def setUp(self):
+        self.program = paddle.static.Program()
+        self.executor = paddle.static.Executor(self.place)
+        with paddle.static.program_guard(self.program):
+            probs = paddle.static.data(
+                'probs', self.probs.shape, self.probs.dtype
+            )
+
+            self._paddle_geometric = geometric.Geometric(probs)
+            self.feeds = {'probs': self.probs, 'value': self.value}
+
+    def test_pmf(self):
+        with paddle.static.program_guard(self.program):
+            [pmf] = self.executor.run(
+                self.program,
+                feed=self.feeds,
+                fetch_list=[self._paddle_geometric.pmf(self.value)],
+            )
+            np.testing.assert_allclose(
+                pmf,
+                scipy.stats.geom.pmf(self.value, self.probs, loc=-1),
+                rtol=RTOL.get(str(self.probs.dtype)),
+                atol=ATOL.get(str(self.probs.dtype)),
+            )
+
+    def test_log_pmf(self):
+        with paddle.static.program_guard(self.program):
+            [log_pmf] = self.executor.run(
+                self.program,
+                feed=self.feeds,
+                fetch_list=[self._paddle_geometric.log_pmf(self.value)],
+            )
+            np.testing.assert_allclose(
+                log_pmf,
+                scipy.stats.geom.logpmf(self.value, self.probs, loc=-1),
+                rtol=RTOL.get(str(self.probs.dtype)),
+                atol=ATOL.get(str(self.probs.dtype)),
+            )
+
+    def test_cdf(self):
+        with paddle.static.program_guard(self.program):
+            [cdf] = self.executor.run(
+                self.program,
+                feed=self.feeds,
+                fetch_list=[self._paddle_geometric.cdf(self.value)],
+            )
+            np.testing.assert_allclose(
+                cdf,
+                scipy.stats.geom.cdf(self.value, self.probs, loc=-1),
+                rtol=RTOL.get(str(self.probs.dtype)),
+                atol=ATOL.get(str(self.probs.dtype)),
+            )
+
+    def test_pmf_error(self):
+        self.assertRaises(TypeError, self._paddle_geometric.pmf, [1, 2])
+
+    def test_log_pmf_error(self):
+        self.assertRaises(TypeError, self._paddle_geometric.log_pmf, [1, 2])
+
+    def test_cdf_error(self):
+        self.assertRaises(TypeError, self._paddle_geometric.cdf, [1, 2])
+
+
+@place(DEVICES)
+@parameterize_cls(
+    (TEST_CASE_NAME, 'probs1', 'probs2'),
+    [
+        (
+            'one-dim',
+            xrand(
+                (2,),
+                dtype='float32',
+                min=np.finfo(dtype='float32').tiny,
+                max=1.0,
+            ),
+            xrand(
+                (2,),
+                dtype='float32',
+                min=np.finfo(dtype='float32').tiny,
+                max=1.0,
+            ),
+        ),
+        (
+            'multi-dim',
+            xrand(
+                (2, 2),
+                dtype='float32',
+                min=np.finfo(dtype='float32').tiny,
+                max=1.0,
+            ),
+            xrand(
+                (2, 2),
+                dtype='float32',
+                min=np.finfo(dtype='float32').tiny,
+                max=1.0,
+            ),
+        ),
+    ],
+)
+class TestGeometricKL(unittest.TestCase):
+    def setUp(self):
+        paddle.enable_static()
+        self.program_p = paddle.static.Program()
+        self.program_q = paddle.static.Program()
+        self.executor = paddle.static.Executor(self.place)
+        with paddle.static.program_guard(self.program_p, self.program_q):
+            probs_p = paddle.static.data(
+                'probs1', self.probs1.shape, self.probs1.dtype
+            )
+            probs_q = paddle.static.data(
+                'probs2', self.probs2.shape, self.probs2.dtype
+            )
+
+            self._paddle_geomP = geometric.Geometric(probs_p)
+            self._paddle_geomQ = geometric.Geometric(probs_q)
+
+            self.feeds = {
+                'probs1': self.probs1,
+                'probs2': self.probs2,
+            }
+
+    def test_kl_divergence(self):
+        with paddle.static.program_guard(self.program_p, self.program_q):
+            self.executor.run(self.program_q)
+            [kl_diver] = self.executor.run(
+                self.program_p,
+                feed=self.feeds,
+                fetch_list=[
+                    self._paddle_geomP.kl_divergence(self._paddle_geomQ)
+                ],
+            )
+            np.testing.assert_allclose(
+                kl_diver,
+                self._kl(),
+                rtol=RTOL.get(str(self.probs1.dtype)),
+                atol=ATOL.get(str(self.probs1.dtype)),
+            )
+
+    def test_kl1_error(self):
+        self.assertRaises(
+            TypeError,
+            self._paddle_geomP.kl_divergence,
+            paddle.distribution.beta.Beta,
+        )
+
+    def test_kl2_error(self):
+        self.assertRaises(
+            TypeError,
+            self._paddle_geomQ.kl_divergence,
+            paddle.distribution.beta.Beta,
+        )
+
+    def _kl(self):
+        return self.probs1 * np.log(self.probs1 / self.probs2) + (
+            1.0 - self.probs1
+        ) * np.log((1.0 - self.probs1) / (1.0 - self.probs2))
+
+
+if __name__ == '__main__':
+    unittest.main()

test/distribution/test_distribution_gumbel.py

@@ -0,0 +1,184 @@
+# Copyright (c) 2022 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
+#
+#     http://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
+import parameterize
+import scipy.stats
+from distribution import config
+
+import paddle
+from paddle.distribution.gumbel import Gumbel
+
+
+@parameterize.place(config.DEVICES)
+@parameterize.parameterize_cls(
+    (parameterize.TEST_CASE_NAME, 'loc', 'scale'),
+    [
+        ('one-dim', parameterize.xrand((4,)), parameterize.xrand((4,))),
+        ('multi-dim', parameterize.xrand((5, 3)), parameterize.xrand((5, 3))),
+    ],
+)
+class TestGumbel(unittest.TestCase):
+    def setUp(self):
+        self._dist = Gumbel(
+            loc=paddle.to_tensor(self.loc), scale=paddle.to_tensor(self.scale)
+        )
+
+    def test_mean(self):
+        mean = self._dist.mean
+        self.assertEqual(mean.numpy().dtype, self._np_mean().dtype)
+        np.testing.assert_allclose(
+            mean,
+            self._np_mean(),
+            rtol=config.RTOL.get(str(self.scale.dtype)),
+            atol=config.ATOL.get(str(self.scale.dtype)),
+        )
+
+    def test_variance(self):
+        var = self._dist.variance
+        self.assertEqual(var.numpy().dtype, self._np_variance().dtype)
+        np.testing.assert_allclose(
+            var,
+            self._np_variance(),
+            rtol=config.RTOL.get(str(self.scale.dtype)),
+            atol=config.ATOL.get(str(self.scale.dtype)),
+        )
+
+    def test_stddev(self):
+        stddev = self._dist.stddev
+        self.assertEqual(stddev.numpy().dtype, self._np_stddev().dtype)
+        np.testing.assert_allclose(
+            stddev,
+            self._np_stddev(),
+            rtol=config.RTOL.get(str(self.scale.dtype)),
+            atol=config.ATOL.get(str(self.scale.dtype)),
+        )
+
+    def test_entropy(self):
+        entropy = self._dist.entropy()
+        self.assertEqual(entropy.numpy().dtype, self._np_entropy().dtype)
+        np.testing.assert_allclose(
+            entropy,
+            self._np_entropy(),
+            rtol=config.RTOL.get(str(self.scale.dtype)),
+            atol=config.ATOL.get(str(self.scale.dtype)),
+        )
+
+    def test_sample(self):
+        sample_shape = [10000]
+        samples = self._dist.sample(sample_shape)
+        sample_values = samples.numpy()
+        self.assertEqual(sample_values.dtype, self.scale.dtype)
+
+        np.testing.assert_allclose(
+            sample_values.mean(axis=0),
+            scipy.stats.gumbel_r.mean(self.loc, scale=self.scale),
+            rtol=0.1,
+            atol=config.ATOL.get(str(self.loc.dtype)),
+        )
+        np.testing.assert_allclose(
+            sample_values.var(axis=0),
+            scipy.stats.gumbel_r.var(self.loc, scale=self.scale),
+            rtol=0.1,
+            atol=config.ATOL.get(str(self.loc.dtype)),
+        )
+
+    def test_rsample(self):
+        sample_shape = [10000]
+        samples = self._dist.rsample(sample_shape)
+        sample_values = samples.numpy()
+        self.assertEqual(sample_values.dtype, self.scale.dtype)
+
+        np.testing.assert_allclose(
+            sample_values.mean(axis=0),
+            scipy.stats.gumbel_r.mean(self.loc, scale=self.scale),
+            rtol=0.1,
+            atol=config.ATOL.get(str(self.loc.dtype)),
+        )
+        np.testing.assert_allclose(
+            sample_values.var(axis=0),
+            scipy.stats.gumbel_r.var(self.loc, scale=self.scale),
+            rtol=0.1,
+            atol=config.ATOL.get(str(self.loc.dtype)),
+        )
+
+    def _np_mean(self):
+        return self.loc + self.scale * np.euler_gamma
+
+    def _np_stddev(self):
+        return np.sqrt(self._np_variance())
+
+    def _np_variance(self):
+        return np.divide(
+            np.multiply(np.power(self.scale, 2), np.power(np.pi, 2)), 6
+        )
+
+    def _np_entropy(self):
+        return np.log(self.scale) + 1 + np.euler_gamma
+
+
+@parameterize.place(config.DEVICES)
+@parameterize.parameterize_cls(
+    (parameterize.TEST_CASE_NAME, 'loc', 'scale', 'value'),
+    [
+        (
+            'value-float',
+            np.array([0.1, 0.4]),
+            np.array([1.0, 4.0]),
+            np.array([3.0, 7.0]),
+        ),
+        ('value-int', np.array([0.1, 0.4]), np.array([1, 4]), np.array([3, 7])),
+        (
+            'value-multi-dim',
+            np.array([0.1, 0.4]),
+            np.array([1, 4]),
+            np.array([[5.0, 4], [6, 2]]),
+        ),
+    ],
+)
+class TestGumbelPDF(unittest.TestCase):
+    def setUp(self):
+        self._dist = Gumbel(
+            loc=paddle.to_tensor(self.loc), scale=paddle.to_tensor(self.scale)
+        )
+
+    def test_prob(self):
+        np.testing.assert_allclose(
+            self._dist.prob(paddle.to_tensor(self.value)),
+            scipy.stats.gumbel_r.pdf(self.value, self.loc, self.scale),
+            rtol=config.RTOL.get(str(self.loc.dtype)),
+            atol=config.ATOL.get(str(self.loc.dtype)),
+        )
+
+    def test_log_prob(self):
+        np.testing.assert_allclose(
+            self._dist.log_prob(paddle.to_tensor(self.value)),
+            scipy.stats.gumbel_r.logpdf(self.value, self.loc, self.scale),
+            rtol=config.RTOL.get(str(self.loc.dtype)),
+            atol=config.ATOL.get(str(self.loc.dtype)),
+        )
+
+    def test_cdf(self):
+        np.testing.assert_allclose(
+            self._dist.cdf(paddle.to_tensor(self.value)),
+            scipy.stats.gumbel_r.cdf(self.value, self.loc, self.scale),
+            rtol=0.02,
+            atol=config.ATOL.get(str(self.loc.dtype)),
+        )
+
+
+if __name__ == '__main__':
+    unittest.main()

test/distribution/test_distribution_gumbel_static.py

@@ -0,0 +1,201 @@
+# Copyright (c) 2022 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
+#
+#     http://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
+import parameterize
+import scipy.stats
+from distribution import config
+
+import paddle
+from paddle.distribution.gumbel import Gumbel
+
+paddle.enable_static()
+
+
+@parameterize.place(config.DEVICES)
+@parameterize.parameterize_cls(
+    (parameterize.TEST_CASE_NAME, 'loc', 'scale'),
+    [
+        ('one-dim', parameterize.xrand((4,)), parameterize.xrand((4,))),
+        ('multi-dim', parameterize.xrand((5, 3)), parameterize.xrand((5, 3))),
+    ],
+)
+class TestGumbel(unittest.TestCase):
+    def setUp(self):
+        startup_program = paddle.static.Program()
+        main_program = paddle.static.Program()
+        executor = paddle.static.Executor(self.place)
+        with paddle.static.program_guard(main_program, startup_program):
+            loc = paddle.static.data('loc', self.loc.shape, self.loc.dtype)
+            scale = paddle.static.data(
+                'scale', self.scale.shape, self.scale.dtype
+            )
+            self._dist = Gumbel(loc=loc, scale=scale)
+            self.sample_shape = [50000]
+            mean = self._dist.mean
+            var = self._dist.variance
+            stddev = self._dist.stddev
+            entropy = self._dist.entropy()
+            samples = self._dist.sample(self.sample_shape)
+        fetch_list = [mean, var, stddev, entropy, samples]
+        self.feeds = {'loc': self.loc, 'scale': self.scale}
+
+        executor.run(startup_program)
+        [
+            self.mean,
+            self.var,
+            self.stddev,
+            self.entropy,
+            self.samples,
+        ] = executor.run(main_program, feed=self.feeds, fetch_list=fetch_list)
+
+    def test_mean(self):
+        self.assertEqual(str(self.mean.dtype).split('.')[-1], self.scale.dtype)
+        np.testing.assert_allclose(
+            self.mean,
+            self._np_mean(),
+            rtol=config.RTOL.get(str(self.scale.dtype)),
+            atol=config.ATOL.get(str(self.scale.dtype)),
+        )
+
+    def test_variance(self):
+        self.assertEqual(str(self.var.dtype).split('.')[-1], self.scale.dtype)
+        np.testing.assert_allclose(
+            self.var,
+            self._np_variance(),
+            rtol=config.RTOL.get(str(self.scale.dtype)),
+            atol=config.ATOL.get(str(self.scale.dtype)),
+        )
+
+    def test_stddev(self):
+        self.assertEqual(
+            str(self.stddev.dtype).split('.')[-1], self.scale.dtype
+        )
+        np.testing.assert_allclose(
+            self.stddev,
+            self._np_stddev(),
+            rtol=config.RTOL.get(str(self.scale.dtype)),
+            atol=config.ATOL.get(str(self.scale.dtype)),
+        )
+
+    def test_entropy(self):
+        self.assertEqual(
+            str(self.entropy.dtype).split('.')[-1], self.scale.dtype
+        )
+
+    def test_sample(self):
+        self.assertEqual(self.samples.dtype, self.scale.dtype)
+
+        np.testing.assert_allclose(
+            self.samples.mean(axis=0),
+            scipy.stats.gumbel_r.mean(self.loc, scale=self.scale),
+            rtol=0.1,
+            atol=config.ATOL.get(str(self.scale.dtype)),
+        )
+        np.testing.assert_allclose(
+            self.samples.var(axis=0),
+            scipy.stats.gumbel_r.var(self.loc, scale=self.scale),
+            rtol=0.1,
+            atol=config.ATOL.get(str(self.scale.dtype)),
+        )
+
+    def _np_mean(self):
+        return self.loc + self.scale * np.euler_gamma
+
+    def _np_stddev(self):
+        return np.sqrt(self._np_variance())
+
+    def _np_variance(self):
+        return np.divide(
+            np.multiply(np.power(self.scale, 2), np.power(np.pi, 2)), 6
+        )
+
+    def _np_entropy(self):
+        return np.log(self.scale) + 1 + np.euler_gamma
+
+
+@parameterize.place(config.DEVICES)
+@parameterize.parameterize_cls(
+    (parameterize.TEST_CASE_NAME, 'loc', 'scale', 'value'),
+    [
+        (
+            'value-float',
+            np.array([0.1, 0.4]),
+            np.array([1.0, 4.0]),
+            np.array([3.0, 7.0]),
+        ),
+        ('value-int', np.array([0.1, 0.4]), np.array([1, 4]), np.array([3, 7])),
+        (
+            'value-multi-dim',
+            np.array([0.1, 0.4]),
+            np.array([1, 4]),
+            np.array([[5.0, 4], [6, 2]]),
+        ),
+    ],
+)
+class TestGumbelPDF(unittest.TestCase):
+    def setUp(self):
+        startup_program = paddle.static.Program()
+        main_program = paddle.static.Program()
+        executor = paddle.static.Executor(self.place)
+
+        with paddle.static.program_guard(main_program, startup_program):
+            loc = paddle.static.data('loc', self.loc.shape, self.loc.dtype)
+            scale = paddle.static.data(
+                'scale', self.scale.shape, self.scale.dtype
+            )
+            value = paddle.static.data(
+                'value', self.value.shape, self.value.dtype
+            )
+            self._dist = Gumbel(loc=loc, scale=scale)
+            prob = self._dist.prob(value)
+            log_prob = self._dist.log_prob(value)
+            cdf = self._dist.cdf(value)
+        fetch_list = [prob, log_prob, cdf]
+        self.feeds = {'loc': self.loc, 'scale': self.scale, 'value': self.value}
+
+        executor.run(startup_program)
+        [self.prob, self.log_prob, self.cdf] = executor.run(
+            main_program, feed=self.feeds, fetch_list=fetch_list
+        )
+
+    def test_prob(self):
+        np.testing.assert_allclose(
+            self.prob,
+            scipy.stats.gumbel_r.pdf(self.value, self.loc, self.scale),
+            rtol=config.RTOL.get(str(self.loc.dtype)),
+            atol=config.ATOL.get(str(self.loc.dtype)),
+        )
+
+    def test_log_prob(self):
+        np.testing.assert_allclose(
+            self.log_prob,
+            scipy.stats.gumbel_r.logpdf(self.value, self.loc, self.scale),
+            rtol=config.RTOL.get(str(self.loc.dtype)),
+            atol=config.ATOL.get(str(self.loc.dtype)),
+        )
+
+    def test_cdf(self):
+        np.testing.assert_allclose(
+            self.cdf,
+            scipy.stats.gumbel_r.cdf(self.value, self.loc, self.scale),
+            rtol=0.3,
+            atol=config.ATOL.get(str(self.loc.dtype)),
+        )
+
+
+if __name__ == '__main__':
+    unittest.main()

test/distribution/test_distribution_independent.py

@@ -0,0 +1,119 @@
+# Copyright (c) 2021 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
+#
+#     http://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
+import parameterize as param
+from distribution import config
+
+import paddle
+
+np.random.seed(2022)
+
+
+@param.place(config.DEVICES)
+@param.param_cls(
+    (param.TEST_CASE_NAME, 'base', 'reinterpreted_batch_rank'),
+    [
+        (
+            'base_beta',
+            paddle.distribution.Beta(paddle.rand([1, 2]), paddle.rand([1, 2])),
+            1,
+        )
+    ],
+)
+class TestIndependent(unittest.TestCase):
+    def setUp(self):
+        self._t = paddle.distribution.Independent(
+            self.base, self.reinterpreted_batch_rank
+        )
+
+    def test_mean(self):
+        np.testing.assert_allclose(
+            self.base.mean,
+            self._t.mean,
+            rtol=config.RTOL.get(str(self.base.alpha.numpy().dtype)),
+            atol=config.ATOL.get(str(self.base.alpha.numpy().dtype)),
+        )
+
+    def test_variance(self):
+        np.testing.assert_allclose(
+            self.base.variance,
+            self._t.variance,
+            rtol=config.RTOL.get(str(self.base.alpha.numpy().dtype)),
+            atol=config.ATOL.get(str(self.base.alpha.numpy().dtype)),
+        )
+
+    def test_entropy(self):
+        np.testing.assert_allclose(
+            self._np_sum_rightmost(
+                self.base.entropy().numpy(), self.reinterpreted_batch_rank
+            ),
+            self._t.entropy(),
+            rtol=config.RTOL.get(str(self.base.alpha.numpy().dtype)),
+            atol=config.ATOL.get(str(self.base.alpha.numpy().dtype)),
+        )
+
+    def _np_sum_rightmost(self, value, n):
+        return np.sum(value, tuple(range(-n, 0))) if n > 0 else value
+
+    def test_log_prob(self):
+        value = np.random.rand(1)
+        np.testing.assert_allclose(
+            self._np_sum_rightmost(
+                self.base.log_prob(paddle.to_tensor(value)).numpy(),
+                self.reinterpreted_batch_rank,
+            ),
+            self._t.log_prob(paddle.to_tensor(value)).numpy(),
+            rtol=config.RTOL.get(str(self.base.alpha.numpy().dtype)),
+            atol=config.ATOL.get(str(self.base.alpha.numpy().dtype)),
+        )
+
+    # TODO(cxxly): Add Kolmogorov-Smirnov test for sample result.
+    def test_sample(self):
+        shape = (5, 10, 8)
+        expected_shape = (5, 10, 8, 1, 2)
+        data = self._t.sample(shape)
+        self.assertEqual(tuple(data.shape), expected_shape)
+        self.assertEqual(data.dtype, self.base.alpha.dtype)
+
+
+@param.place(config.DEVICES)
+@param.param_cls(
+    (
+        param.TEST_CASE_NAME,
+        'base',
+        'reinterpreted_batch_rank',
+        'expected_exception',
+    ),
+    [
+        ('base_not_transform', '', 1, TypeError),
+        (
+            'rank_less_than_zero',
+            paddle.distribution.Transform(),
+            -1,
+            ValueError,
+        ),
+    ],
+)
+class TestIndependentException(unittest.TestCase):
+    def test_init(self):
+        with self.assertRaises(self.expected_exception):
+            paddle.distribution.IndependentTransform(
+                self.base, self.reinterpreted_batch_rank
+            )
+
+
+if __name__ == '__main__':
+    unittest.main()

test/distribution/test_distribution_independent_static.py

@@ -0,0 +1,133 @@
+# Copyright (c) 2021 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
+#
+#     http://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
+import parameterize as param
+from distribution import config
+
+import paddle
+
+np.random.seed(2022)
+paddle.enable_static()
+
+
+@param.place(config.DEVICES)
+@param.param_cls(
+    (param.TEST_CASE_NAME, 'base', 'reinterpreted_batch_rank', 'alpha', 'beta'),
+    [
+        (
+            'base_beta',
+            paddle.distribution.Beta,
+            1,
+            np.random.rand(1, 2),
+            np.random.rand(1, 2),
+        )
+    ],
+)
+class TestIndependent(unittest.TestCase):
+    def setUp(self):
+        value = np.random.rand(1)
+        self.dtype = value.dtype
+        exe = paddle.static.Executor()
+        sp = paddle.static.Program()
+        mp = paddle.static.Program()
+        with paddle.static.program_guard(mp, sp):
+            alpha = paddle.static.data(
+                'alpha', self.alpha.shape, self.alpha.dtype
+            )
+            beta = paddle.static.data('beta', self.beta.shape, self.beta.dtype)
+            self.base = self.base(alpha, beta)
+            t = paddle.distribution.Independent(
+                self.base, self.reinterpreted_batch_rank
+            )
+            mean = t.mean
+            variance = t.variance
+            entropy = t.entropy()
+            static_value = paddle.static.data('value', value.shape, value.dtype)
+            log_prob = t.log_prob(static_value)
+
+            base_mean = self.base.mean
+            base_variance = self.base.variance
+            base_entropy = self.base.entropy()
+            base_log_prob = self.base.log_prob(static_value)
+
+        fetch_list = [
+            mean,
+            variance,
+            entropy,
+            log_prob,
+            base_mean,
+            base_variance,
+            base_entropy,
+            base_log_prob,
+        ]
+        exe.run(sp)
+        [
+            self.mean,
+            self.variance,
+            self.entropy,
+            self.log_prob,
+            self.base_mean,
+            self.base_variance,
+            self.base_entropy,
+            self.base_log_prob,
+        ] = exe.run(
+            mp,
+            feed={'value': value, 'alpha': self.alpha, 'beta': self.beta},
+            fetch_list=fetch_list,
+        )
+
+    def test_mean(self):
+        np.testing.assert_allclose(
+            self.mean,
+            self.base_mean,
+            rtol=config.RTOL.get(str(self.dtype)),
+            atol=config.ATOL.get(str(self.dtype)),
+        )
+
+    def test_variance(self):
+        np.testing.assert_allclose(
+            self.variance,
+            self.base_variance,
+            rtol=config.RTOL.get(str(self.dtype)),
+            atol=config.ATOL.get(str(self.dtype)),
+        )
+
+    def test_entropy(self):
+        np.testing.assert_allclose(
+            self._np_sum_rightmost(
+                self.base_entropy, self.reinterpreted_batch_rank
+            ),
+            self.entropy,
+            rtol=config.RTOL.get(str(self.dtype)),
+            atol=config.ATOL.get(str(self.dtype)),
+        )
+
+    def _np_sum_rightmost(self, value, n):
+        return np.sum(value, tuple(range(-n, 0))) if n > 0 else value
+
+    def test_log_prob(self):
+        np.testing.assert_allclose(
+            self._np_sum_rightmost(
+                self.base_log_prob, self.reinterpreted_batch_rank
+            ),
+            self.log_prob,
+            rtol=config.RTOL.get(str(self.dtype)),
+            atol=config.ATOL.get(str(self.dtype)),
+        )
+
+
+if __name__ == '__main__':
+    unittest.main()

test/distribution/test_distribution_laplace.py

@@ -0,0 +1,242 @@
+# Copyright (c) 2022 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
+#
+#     http://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
+import parameterize
+import scipy.stats
+from distribution import config
+
+import paddle
+
+
+@parameterize.place(config.DEVICES)
+@parameterize.parameterize_cls(
+    (parameterize.TEST_CASE_NAME, 'loc', 'scale'),
+    [
+        ('one-dim', parameterize.xrand((2,)), parameterize.xrand((2,))),
+        ('multi-dim', parameterize.xrand((5, 5)), parameterize.xrand((5, 5))),
+    ],
+)
+class TestLaplace(unittest.TestCase):
+    def setUp(self):
+        self._dist = paddle.distribution.Laplace(
+            loc=paddle.to_tensor(self.loc), scale=paddle.to_tensor(self.scale)
+        )
+
+    def test_mean(self):
+        mean = self._dist.mean
+        self.assertEqual(mean.numpy().dtype, self.scale.dtype)
+        np.testing.assert_allclose(
+            mean,
+            self._np_mean(),
+            rtol=config.RTOL.get(str(self.scale.dtype)),
+            atol=config.ATOL.get(str(self.scale.dtype)),
+        )
+
+    def test_variance(self):
+        var = self._dist.variance
+        self.assertEqual(var.numpy().dtype, self.scale.dtype)
+        np.testing.assert_allclose(
+            var,
+            self._np_variance(),
+            rtol=config.RTOL.get(str(self.scale.dtype)),
+            atol=config.ATOL.get(str(self.scale.dtype)),
+        )
+
+    def test_stddev(self):
+        stddev = self._dist.stddev
+        self.assertEqual(stddev.numpy().dtype, self.scale.dtype)
+        np.testing.assert_allclose(
+            stddev,
+            self._np_stddev(),
+            rtol=config.RTOL.get(str(self.scale.dtype)),
+            atol=config.ATOL.get(str(self.scale.dtype)),
+        )
+
+    def test_entropy(self):
+        entropy = self._dist.entropy()
+        self.assertEqual(entropy.numpy().dtype, self.scale.dtype)
+
+    def test_sample(self):
+        sample_shape = (50000,)
+        samples = self._dist.sample(sample_shape)
+        sample_values = samples.numpy()
+
+        self.assertEqual(samples.numpy().dtype, self.scale.dtype)
+        self.assertEqual(
+            tuple(samples.shape), tuple(self._dist._extend_shape(sample_shape))
+        )
+
+        self.assertEqual(samples.shape, list(sample_shape + self.loc.shape))
+        self.assertEqual(sample_values.shape, sample_shape + self.loc.shape)
+
+        np.testing.assert_allclose(
+            sample_values.mean(axis=0),
+            scipy.stats.laplace.mean(self.loc, scale=self.scale),
+            rtol=0.2,
+            atol=0.0,
+        )
+        np.testing.assert_allclose(
+            sample_values.var(axis=0),
+            scipy.stats.laplace.var(self.loc, scale=self.scale),
+            rtol=0.1,
+            atol=0.0,
+        )
+
+    def _np_mean(self):
+        return self.loc
+
+    def _np_stddev(self):
+        return (2**0.5) * self.scale
+
+    def _np_variance(self):
+        stddev = (2**0.5) * self.scale
+        return np.power(stddev, 2)
+
+    def _np_entropy(self):
+        return scipy.stats.laplace.entropy(loc=self.loc, scale=self.scale)
+
+
+@parameterize.place(config.DEVICES)
+@parameterize.parameterize_cls(
+    (parameterize.TEST_CASE_NAME, 'loc', 'scale'),
+    [
+        ('float', 1.0, 2.0),
+        ('int', 3, 4),
+    ],
+)
+class TestLaplaceKS(unittest.TestCase):
+    def setUp(self):
+        self._dist = paddle.distribution.Laplace(loc=self.loc, scale=self.scale)
+
+    def test_sample(self):
+        sample_shape = (20000,)
+        samples = self._dist.sample(sample_shape)
+        sample_values = samples.numpy()
+        self.assertTrue(self._kstest(self.loc, self.scale, sample_values))
+
+    def _kstest(self, loc, scale, samples):
+        # Uses the Kolmogorov-Smirnov test for goodness of fit.
+        ks, p_value = scipy.stats.kstest(
+            samples, scipy.stats.laplace(loc, scale=scale).cdf
+        )
+        return ks < 0.02
+
+
+@parameterize.place(config.DEVICES)
+@parameterize.parameterize_cls(
+    (parameterize.TEST_CASE_NAME, 'loc', 'scale', 'value'),
+    [
+        (
+            'value-float',
+            np.array([0.2, 0.3]),
+            np.array([2.0, 3.0]),
+            np.array([2.0, 5.0]),
+        ),
+        (
+            'value-int',
+            np.array([0.2, 0.3]),
+            np.array([2.0, 3.0]),
+            np.array([2, 5]),
+        ),
+        (
+            'value-multi-dim',
+            np.array([0.2, 0.3]),
+            np.array([2.0, 3.0]),
+            np.array([[4.0, 6], [8, 2]]),
+        ),
+    ],
+)
+class TestLaplacePDF(unittest.TestCase):
+    def setUp(self):
+        self._dist = paddle.distribution.Laplace(
+            loc=paddle.to_tensor(self.loc), scale=paddle.to_tensor(self.scale)
+        )
+
+    def test_prob(self):
+        np.testing.assert_allclose(
+            self._dist.prob(paddle.to_tensor(self.value)),
+            scipy.stats.laplace.pdf(self.value, self.loc, self.scale),
+            rtol=config.RTOL.get(str(self.loc.dtype)),
+            atol=config.ATOL.get(str(self.loc.dtype)),
+        )
+
+    def test_log_prob(self):
+        np.testing.assert_allclose(
+            self._dist.log_prob(paddle.to_tensor(self.value)),
+            scipy.stats.laplace.logpdf(self.value, self.loc, self.scale),
+            rtol=config.RTOL.get(str(self.loc.dtype)),
+            atol=config.ATOL.get(str(self.loc.dtype)),
+        )
+
+    def test_cdf(self):
+        np.testing.assert_allclose(
+            self._dist.cdf(paddle.to_tensor(self.value)),
+            scipy.stats.laplace.cdf(self.value, self.loc, self.scale),
+            rtol=config.RTOL.get(str(self.loc.dtype)),
+            atol=config.ATOL.get(str(self.loc.dtype)),
+        )
+
+    def test_icdf(self):
+        np.testing.assert_allclose(
+            self._dist.icdf(paddle.to_tensor(self.value)),
+            scipy.stats.laplace.ppf(self.value, self.loc, self.scale),
+            rtol=config.RTOL.get(str(self.loc.dtype)),
+            atol=config.ATOL.get(str(self.loc.dtype)),
+        )
+
+
+@parameterize.place(config.DEVICES)
+@parameterize.parameterize_cls(
+    (parameterize.TEST_CASE_NAME, 'loc1', 'scale1', 'loc2', 'scale2'),
+    [
+        (
+            'kl',
+            np.array([0.0]),
+            np.array([1.0]),
+            np.array([1.0]),
+            np.array([0.5]),
+        )
+    ],
+)
+class TestLaplaceAndLaplaceKL(unittest.TestCase):
+    def setUp(self):
+        self._dist_1 = paddle.distribution.Laplace(
+            loc=paddle.to_tensor(self.loc1), scale=paddle.to_tensor(self.scale1)
+        )
+        self._dist_2 = paddle.distribution.Laplace(
+            loc=paddle.to_tensor(self.loc2), scale=paddle.to_tensor(self.scale2)
+        )
+
+    def test_kl_divergence(self):
+        np.testing.assert_allclose(
+            paddle.distribution.kl_divergence(self._dist_1, self._dist_2),
+            self._np_kl(),
+            atol=0,
+            rtol=0.50,
+        )
+
+    def _np_kl(self):
+        x = np.linspace(
+            scipy.stats.laplace.ppf(0.01), scipy.stats.laplace.ppf(0.99), 1000
+        )
+        d1 = scipy.stats.laplace.pdf(x, loc=0.0, scale=1.0)
+        d2 = scipy.stats.laplace.pdf(x, loc=1.0, scale=0.5)
+        return scipy.stats.entropy(d1, d2)
+
+
+if __name__ == '__main__':
+    unittest.main()

test/distribution/test_distribution_laplace_static.py

@@ -0,0 +1,343 @@
+# Copyright (c) 2022 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
+#
+#     http://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
+import parameterize
+import scipy.stats
+from distribution import config
+
+import paddle
+
+paddle.enable_static()
+
+
+@parameterize.place(config.DEVICES)
+@parameterize.parameterize_cls(
+    (parameterize.TEST_CASE_NAME, 'loc', 'scale'),
+    [
+        ('one-dim', parameterize.xrand((2,)), parameterize.xrand((2,))),
+        ('multi-dim', parameterize.xrand((5, 5)), parameterize.xrand((5, 5))),
+    ],
+    test_pir=True,
+)
+class TestLaplace(unittest.TestCase):
+    def build_program(self):
+        startup_program = paddle.static.Program()
+        main_program = paddle.static.Program()
+        executor = paddle.static.Executor(self.place)
+        with paddle.static.program_guard(main_program, startup_program):
+            loc = paddle.static.data('loc', self.loc.shape, self.loc.dtype)
+            scale = paddle.static.data(
+                'scale', self.scale.shape, self.scale.dtype
+            )
+            self._dist = paddle.distribution.Laplace(loc=loc, scale=scale)
+            self.sample_shape = (50000,)
+            mean = self._dist.mean
+            var = self._dist.variance
+            stddev = self._dist.stddev
+            entropy = self._dist.entropy()
+            samples = self._dist.sample(self.sample_shape)
+        fetch_list = [mean, var, stddev, entropy, samples]
+        self.feeds = {'loc': self.loc, 'scale': self.scale}
+
+        executor.run(startup_program)
+        [
+            self.mean,
+            self.var,
+            self.stddev,
+            self.entropy,
+            self.samples,
+        ] = executor.run(main_program, feed=self.feeds, fetch_list=fetch_list)
+
+    def setUp(self):
+        if self.test_pir:
+            with paddle.pir_utils.IrGuard():
+                self.build_program()
+        else:
+            self.build_program()
+
+    def test_mean(self):
+        self.assertEqual(str(self.mean.dtype).split('.')[-1], self.scale.dtype)
+        np.testing.assert_allclose(
+            self.mean,
+            self._np_mean(),
+            rtol=config.RTOL.get(str(self.scale.dtype)),
+            atol=config.ATOL.get(str(self.scale.dtype)),
+        )
+
+    def test_variance(self):
+        self.assertEqual(str(self.var.dtype).split('.')[-1], self.scale.dtype)
+        np.testing.assert_allclose(
+            self.var,
+            self._np_variance(),
+            rtol=config.RTOL.get(str(self.scale.dtype)),
+            atol=config.ATOL.get(str(self.scale.dtype)),
+        )
+
+    def test_stddev(self):
+        self.assertEqual(
+            str(self.stddev.dtype).split('.')[-1], self.scale.dtype
+        )
+        np.testing.assert_allclose(
+            self.stddev,
+            self._np_stddev(),
+            rtol=config.RTOL.get(str(self.scale.dtype)),
+            atol=config.ATOL.get(str(self.scale.dtype)),
+        )
+
+    def test_entropy(self):
+        self.assertEqual(
+            str(self.entropy.dtype).split('.')[-1], self.scale.dtype
+        )
+
+    def test_sample(self):
+        self.assertEqual(self.samples.dtype, self.scale.dtype)
+        self.assertEqual(
+            tuple(self.samples.shape),
+            tuple(self._dist._extend_shape(self.sample_shape)),
+        )
+
+        self.assertEqual(self.samples.shape, self.sample_shape + self.loc.shape)
+        self.assertEqual(self.samples.shape, self.sample_shape + self.loc.shape)
+
+        np.testing.assert_allclose(
+            self.samples.mean(axis=0),
+            scipy.stats.laplace.mean(self.loc, scale=self.scale),
+            rtol=0.2,
+            atol=0.0,
+        )
+        np.testing.assert_allclose(
+            self.samples.var(axis=0),
+            scipy.stats.laplace.var(self.loc, scale=self.scale),
+            rtol=0.1,
+            atol=0.0,
+        )
+
+    def _np_mean(self):
+        return self.loc
+
+    def _np_stddev(self):
+        return (2**0.5) * self.scale
+
+    def _np_variance(self):
+        stddev = (2**0.5) * self.scale
+        return np.power(stddev, 2)
+
+    def _np_entropy(self):
+        return scipy.stats.laplace.entropy(loc=self.loc, scale=self.scale)
+
+
+@parameterize.place(config.DEVICES)
+@parameterize.parameterize_cls(
+    (parameterize.TEST_CASE_NAME, 'loc', 'scale', 'value'),
+    [
+        (
+            'value-float',
+            np.array([0.2, 0.3]),
+            np.array([2.0, 3.0]),
+            np.array([2.0, 5.0]),
+        ),
+        (
+            'value-int',
+            np.array([0.2, 0.3]),
+            np.array([2.0, 3.0]),
+            np.array([2, 5]),
+        ),
+        (
+            'value-multi-dim',
+            np.array([0.2, 0.3]),
+            np.array([2.0, 3.0]),
+            np.array([[4.0, 6], [8, 2]]),
+        ),
+    ],
+    test_pir=True,
+)
+class TestLaplacePDF(unittest.TestCase):
+    def build_program(self):
+        startup_program = paddle.static.Program()
+        main_program = paddle.static.Program()
+        executor = paddle.static.Executor(self.place)
+
+        with paddle.static.program_guard(main_program, startup_program):
+            loc = paddle.static.data('loc', self.loc.shape, self.loc.dtype)
+            scale = paddle.static.data(
+                'scale', self.scale.shape, self.scale.dtype
+            )
+            value = paddle.static.data(
+                'value', self.value.shape, self.value.dtype
+            )
+            self._dist = paddle.distribution.Laplace(loc=loc, scale=scale)
+            prob = self._dist.prob(value)
+            log_prob = self._dist.log_prob(value)
+            cdf = self._dist.cdf(value)
+            icdf = self._dist.icdf(value)
+        fetch_list = [prob, log_prob, cdf, icdf]
+        self.feeds = {'loc': self.loc, 'scale': self.scale, 'value': self.value}
+
+        executor.run(startup_program)
+        [self.prob, self.log_prob, self.cdf, self.icdf] = executor.run(
+            main_program, feed=self.feeds, fetch_list=fetch_list
+        )
+
+    def setUp(self):
+        if self.test_pir:
+            with paddle.pir_utils.IrGuard():
+                self.build_program()
+        else:
+            self.build_program()
+
+    def test_prob(self):
+        np.testing.assert_allclose(
+            self.prob,
+            scipy.stats.laplace.pdf(self.value, self.loc, self.scale),
+            rtol=config.RTOL.get(str(self.loc.dtype)),
+            atol=config.ATOL.get(str(self.loc.dtype)),
+        )
+
+    def test_log_prob(self):
+        np.testing.assert_allclose(
+            self.log_prob,
+            scipy.stats.laplace.logpdf(self.value, self.loc, self.scale),
+            rtol=config.RTOL.get(str(self.loc.dtype)),
+            atol=config.ATOL.get(str(self.loc.dtype)),
+        )
+
+    def test_cdf(self):
+        np.testing.assert_allclose(
+            self.cdf,
+            scipy.stats.laplace.cdf(self.value, self.loc, self.scale),
+            rtol=config.RTOL.get(str(self.loc.dtype)),
+            atol=config.ATOL.get(str(self.loc.dtype)),
+        )
+
+    def test_icdf(self):
+        np.testing.assert_allclose(
+            self.icdf,
+            scipy.stats.laplace.ppf(self.value, self.loc, self.scale),
+            rtol=config.RTOL.get(str(self.loc.dtype)),
+            atol=config.ATOL.get(str(self.loc.dtype)),
+        )
+
+
+@parameterize.place(config.DEVICES)
+@parameterize.parameterize_cls(
+    (parameterize.TEST_CASE_NAME, 'loc1', 'scale1', 'loc2', 'scale2'),
+    [
+        (
+            'kl',
+            np.array([0.0]),
+            np.array([1.0]),
+            np.array([1.0]),
+            np.array([0.5]),
+        )
+    ],
+    test_pir=True,
+)
+class TestLaplaceAndLaplaceKL(unittest.TestCase):
+    def build_program(self):
+        self.mp = paddle.static.Program()
+        self.sp = paddle.static.Program()
+        self.executor = paddle.static.Executor(self.place)
+
+        with paddle.static.program_guard(self.mp, self.sp):
+            loc1 = paddle.static.data('loc1', self.loc1.shape, self.loc1.dtype)
+            scale1 = paddle.static.data(
+                'scale1', self.scale1.shape, self.scale1.dtype
+            )
+            loc2 = paddle.static.data('loc2', self.loc2.shape, self.loc2.dtype)
+            scale2 = paddle.static.data(
+                'scale2', self.scale2.shape, self.scale2.dtype
+            )
+            self._dist_1 = paddle.distribution.Laplace(loc=loc1, scale=scale1)
+            self._dist_2 = paddle.distribution.Laplace(loc=loc2, scale=scale2)
+            self.feeds = {
+                'loc1': self.loc1,
+                'scale1': self.scale1,
+                'loc2': self.loc2,
+                'scale2': self.scale2,
+            }
+
+    def setUp(self):
+        if self.test_pir:
+            with paddle.pir_utils.IrGuard():
+                self.build_program()
+        else:
+            self.build_program()
+
+    def add_kl_divergence(self):
+        with paddle.static.program_guard(self.mp, self.sp):
+            out = paddle.distribution.kl_divergence(self._dist_1, self._dist_2)
+            self.executor.run(self.sp)
+            [out] = self.executor.run(
+                self.mp, feed=self.feeds, fetch_list=[out]
+            )
+            np.testing.assert_allclose(out, self._np_kl(), atol=0, rtol=0.50)
+
+    def test_kl_divergence(self):
+        if self.test_pir:
+            with paddle.pir_utils.IrGuard():
+                self.add_kl_divergence()
+        else:
+            self.add_kl_divergence()
+
+    def _np_kl(self):
+        x = np.linspace(
+            scipy.stats.laplace.ppf(0.01), scipy.stats.laplace.ppf(0.99), 1000
+        )
+        d1 = scipy.stats.laplace.pdf(x, loc=0.0, scale=1.0)
+        d2 = scipy.stats.laplace.pdf(x, loc=1.0, scale=0.5)
+        return scipy.stats.entropy(d1, d2)
+
+
+"""
+# Note: Zero dimension of a Tensor is not supported by static graph mode of paddle;
+# therefore, ks test below cannot be conducted temporarily.
+
+@parameterize.place(config.DEVICES)
+@parameterize.parameterize_cls(
+    (parameterize.TEST_CASE_NAME, 'loc', 'scale', 'sample_shape'), [
+    ('one-dim', np.array(4.0), np.array(3.0), np.array([3000]))])
+class TestLaplaceKS(unittest.TestCase):
+    def setUp(self):
+        self.program = paddle.static.Program()
+        self.executor = paddle.static.Executor(self.place)
+        with paddle.static.program_guard(self.program):
+            loc = paddle.static.data('loc', self.loc.shape,
+                                        self.loc.dtype)
+            scale = paddle.static.data('scale', self.scale.shape,
+                                        self.scale.dtype)
+            self.sample = paddle.static.data('sample_shape', self.sample_shape.shape,
+                                        self.sample_shape.dtype)
+            self._dist = paddle.distribution.Laplace(loc=loc, scale=scale)
+            self.feeds = {'loc': self.loc, 'scale': self.scale, 'sample_shape': self.sample_shape}
+
+    def test_sample(self):
+        with paddle.static.program_guard(self.program):
+            [sample_values] = self.executor.run(self.program,
+                              feed=self.feeds,
+                              fetch_list=self._dist.sample((3000,)))
+            self.assertTrue(self._kstest(self.loc, self.scale, sample_values))
+
+    def _kstest(self, loc, scale, samples):
+        # Uses the Kolmogorov-Smirnov test for goodness of fit.
+        ks, p_value = scipy.stats.kstest(
+            samples,
+            scipy.stats.laplace(loc, scale=scale).cdf)
+        return ks < 0.02
+"""
+
+if __name__ == '__main__':
+    unittest.main()

test/distribution/test_distribution_lkj_cholesky.py

@@ -0,0 +1,188 @@
+# Copyright (c) 2024 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
+#
+#     http://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
+import parameterize
+from distribution import config
+
+import paddle
+from paddle.distribution import lkj_cholesky
+from paddle.distribution.lkj_cholesky import (
+    tril_matrix_to_vec,
+    vec_to_tril_matrix,
+)
+
+np.random.seed(2024)
+paddle.seed(2024)
+
+
+@parameterize.place(config.DEVICES)
+@parameterize.parameterize_cls(
+    (parameterize.TEST_CASE_NAME, 'concentration'),
+    [
+        (
+            'zero-dim',
+            parameterize.xrand(
+                (1,),
+                dtype='float32',
+                max=1.0,
+                min=0,
+            ).reshape([]),
+        ),
+        (
+            'one-dim',
+            parameterize.xrand(
+                (2,),
+                dtype='float32',
+                max=1.0,
+                min=0,
+            ),
+        ),
+        (
+            'one-dim2',
+            parameterize.xrand(
+                (1,),
+                dtype='float32',
+                max=1.0,
+                min=0,
+            ),
+        ),
+    ],
+)
+class TestLKJCholeskyShape(unittest.TestCase):
+    def gen_cases(self):
+        extra_shape = []
+        extra_shape.extend(self.concentration.shape)
+        extra_shape.extend(
+            [self._paddle_lkj_cholesky.dim, self._paddle_lkj_cholesky.dim]
+        )
+        cases = [
+            {
+                'input': (),
+                'expect': tuple(extra_shape),
+            },
+        ]
+        return cases
+
+    def test_onion_sample_shape(self):
+        sample_method = 'onion'
+        self._test_sample_shape_dim(sample_method)
+
+    def test_cvine_sample_shape(self):
+        sample_method = 'cvine'
+        self._test_sample_shape_dim(sample_method)
+
+    def _test_sample_shape_dim(self, sample_method):
+        self._test_sample_shape(2, sample_method)
+
+    def _test_sample_shape(self, dim, sample_method):
+        self._paddle_lkj_cholesky = lkj_cholesky.LKJCholesky(
+            dim, self.concentration, sample_method
+        )
+        cases = self.gen_cases()
+        for case in cases:
+            data = self._paddle_lkj_cholesky.sample(case.get('input'))
+            self.assertTrue(tuple(data.shape) == case.get('expect'))
+
+
+@parameterize.place(config.DEVICES)
+@parameterize.parameterize_cls(
+    (parameterize.TEST_CASE_NAME),
+    [
+        ('test_log_prob'),
+    ],
+)
+class TestLKJCholeskyLogProb(unittest.TestCase):
+    def test_log_prob(self):
+        self.dim = 2
+        self._paddle_lkj_cholesky = lkj_cholesky.LKJCholesky(
+            self.dim, [1.0], 'onion'
+        )
+        self._test_log_prob()
+
+    def _test_log_prob(self):
+        log_probs = []
+        for i in range(2):
+            sample = self._paddle_lkj_cholesky.sample()
+            log_prob = self._paddle_lkj_cholesky.log_prob(sample)
+            sample_tril = tril_matrix_to_vec(sample, diag=-1)
+            # log_abs_det_jacobian
+            logabsdet = []
+            logabsdet.append(self._compute_jacobian(sample_tril)[1])
+            logabsdet = paddle.to_tensor(logabsdet)
+
+            log_probs.append((log_prob - logabsdet).numpy())
+        np.testing.assert_allclose(
+            log_probs[0],
+            log_probs[1],
+            rtol=0.1,
+            atol=config.ATOL.get('float32'),
+        )
+
+    def _tril_cholesky_to_tril_corr(self, x):
+        last_dim = self.dim * (self.dim - 1) // 2
+        x = x.reshape((last_dim,))
+        x = vec_to_tril_matrix(x, self.dim, last_dim, last_dim, (1,), -1)
+        diag = (1 - (x * x).sum(-1)).sqrt().diag_embed()
+        x = x + diag
+        x = x.reshape((self.dim, self.dim))
+        return tril_matrix_to_vec(paddle.matmul(x, x, transpose_y=True), -1)
+
+    def _compute_jacobian(self, x):
+        if x.stop_gradient is not False:
+            x.stop_gradient = False
+        jacobian_matrix = []
+        outputs = self._tril_cholesky_to_tril_corr(x)
+        for i in range(outputs.shape[0]):
+            grad = paddle.grad(
+                outputs=outputs[i], inputs=x, create_graph=False
+            )[0]
+            jacobian_matrix.append(grad)
+        J = paddle.stack(jacobian_matrix, axis=0)
+        logabsdet = paddle.linalg.slogdet(J)
+        return logabsdet
+
+
+@parameterize.place(config.DEVICES)
+@parameterize.parameterize_cls(
+    (parameterize.TEST_CASE_NAME),
+    [
+        ('lkj_cholesky_test_err'),
+    ],
+)
+class LKJCholeskyTestError(unittest.TestCase):
+    @parameterize.parameterize_func(
+        [
+            (1, 1.0, ValueError),  # dim < 2
+            (3.0, 1.0, TypeError),  # dim is float
+            (3, -1.0, ValueError),  # concentration < 0
+        ]
+    )
+    def test_bad_parameter(self, dim, concentration, error):
+        with paddle.base.dygraph.guard(self.place):
+            self.assertRaises(
+                error, lkj_cholesky.LKJCholesky, dim, concentration
+            )
+
+    @parameterize.parameterize_func([(10,)])  # not sequence object sample shape
+    def test_bad_sample_shape(self, shape):
+        with paddle.base.dygraph.guard(self.place):
+            lkj = lkj_cholesky.LKJCholesky(3)
+            self.assertRaises(TypeError, lkj.sample, shape)
+
+
+if __name__ == '__main__':
+    unittest.main()

test/distribution/test_distribution_lkj_cholesky_static.py

@@ -0,0 +1,259 @@
+# Copyright (c) 2024 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
+#
+#     http://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
+import parameterize
+from distribution import config
+
+import paddle
+from paddle.distribution import lkj_cholesky
+from paddle.distribution.lkj_cholesky import (
+    tril_matrix_to_vec,
+    vec_to_tril_matrix,
+)
+
+paddle.enable_static()
+
+np.random.seed(2024)
+paddle.seed(2024)
+
+
+@parameterize.place(config.DEVICES)
+@parameterize.parameterize_cls(
+    (parameterize.TEST_CASE_NAME, 'concentration'),
+    [
+        (
+            'zero-dim',
+            parameterize.xrand(
+                (1,),
+                dtype='float32',
+                max=1.0,
+                min=0,
+            ).reshape([]),
+        ),
+        (
+            'one-dim',
+            parameterize.xrand(
+                (2,),
+                dtype='float32',
+                max=1.0,
+                min=0,
+            ),
+        ),
+        (
+            'one-dim2',
+            parameterize.xrand(
+                (1,),
+                dtype='float64',
+                max=1.0,
+                min=0,
+            ),
+        ),
+    ],
+)
+class TestLKJCholeskyShape(unittest.TestCase):
+    def setUp(self):
+        self.program = paddle.static.Program()
+        self.executor = paddle.static.Executor(self.place)
+        with paddle.static.program_guard(self.program):
+            conc = paddle.static.data(
+                'concentration',
+                self.concentration.shape,
+                self.concentration.dtype,
+            )
+
+            self._paddle_lkj_cholesky_onion = lkj_cholesky.LKJCholesky(
+                2, conc, 'onion'
+            )
+
+            self._paddle_lkj_cholesky_cvine = lkj_cholesky.LKJCholesky(
+                2, conc, 'cvine'
+            )
+
+            self.feeds = {
+                'concentration': self.concentration,
+            }
+
+    def gen_cases(self):
+        extra_shape = []
+        extra_shape.extend(self.concentration.shape)
+        extra_shape.extend(
+            [self._paddle_lkj_cholesky.dim, self._paddle_lkj_cholesky.dim]
+        )
+        cases = [
+            {
+                'input': (),
+                'expect': tuple(extra_shape),
+            },
+            {
+                'input': (2, 2),
+                'expect': (2, 2, *extra_shape),
+            },
+        ]
+        return cases
+
+    def test_onion_sample_shape(self):
+        sample_method = 'onion'
+        self._test_sample_shape(sample_method)
+
+    def test_cvine_sample_shape(self):
+        sample_method = 'cvine'
+        self._test_sample_shape(sample_method)
+
+    def _test_sample_shape(self, sample_method):
+        with paddle.static.program_guard(self.program):
+            if sample_method == 'cvine':
+                self._paddle_lkj_cholesky = self._paddle_lkj_cholesky_cvine
+            else:
+                self._paddle_lkj_cholesky = self._paddle_lkj_cholesky_onion
+
+            cases = self.gen_cases()
+            for case in cases:
+                [data] = self.executor.run(
+                    self.program,
+                    feed=self.feeds,
+                    fetch_list=self._paddle_lkj_cholesky.sample(
+                        case.get('input')
+                    ),
+                )
+            self.assertTrue(tuple(data.shape) == case.get('expect'))
+
+
+@parameterize.place(config.DEVICES)
+@parameterize.parameterize_cls(
+    (parameterize.TEST_CASE_NAME),
+    [
+        ('test_log_prob'),
+    ],
+)
+class TestLKJCholeskyLogProb(unittest.TestCase):
+    def setUp(self):
+        self.program = paddle.static.Program()
+        self.executor = paddle.static.Executor(self.place)
+        with paddle.static.program_guard(self.program):
+            self.concentration = [1.0]
+            self.feeds = {
+                'concentration': self.concentration,
+            }
+
+    def test_log_prob(self):
+        self.dim = 2
+        self._test_log_prob('onion')
+
+    def _test_log_prob(self, sample_method):
+        with paddle.static.program_guard(self.program):
+            self._paddle_lkj_cholesky = lkj_cholesky.LKJCholesky(
+                self.dim, self.concentration, sample_method
+            )
+            log_probs = []
+            for i in range(2):
+                [sample] = self.executor.run(
+                    self.program,
+                    feed=self.feeds,
+                    fetch_list=self._paddle_lkj_cholesky.sample(),
+                )
+                if paddle.framework.use_pir_api():
+                    block = self.program.global_block()
+                    for op in block.ops:
+                        if op.name() == 'pd_op.fetch':
+                            block.remove_op(op)
+                sample = paddle.to_tensor(sample)
+                [log_prob] = self.executor.run(
+                    self.program,
+                    feed=self.feeds,
+                    fetch_list=self._paddle_lkj_cholesky.log_prob(sample),
+                )
+                if paddle.framework.use_pir_api():
+                    block = self.program.global_block()
+                    for op in block.ops:
+                        if op.name() == 'pd_op.fetch':
+                            block.remove_op(op)
+
+                sample_tril = tril_matrix_to_vec(sample, diag=-1)
+
+                # log_abs_det_jacobian
+                logabsdet = []
+
+                [logabsdet_value, J] = self.executor.run(
+                    self.program,
+                    feed={'sample_tril': sample_tril},
+                    fetch_list=self._compute_jacobian(sample_tril),
+                )
+                if paddle.framework.use_pir_api():
+                    block = self.program.global_block()
+                    for op in block.ops:
+                        if op.name() == 'pd_op.fetch':
+                            block.remove_op(op)
+
+                logabsdet.append(logabsdet_value[1])
+
+                log_probs.append(log_prob - logabsdet)
+            max_abs_error = np.max(np.abs(log_probs[0] - log_probs[1]))
+            self.assertAlmostEqual(max_abs_error, 0, places=3)
+
+    def _tril_cholesky_to_tril_corr(self, x, last_dim):
+        x = x.reshape((last_dim,))
+        x = vec_to_tril_matrix(x, self.dim, last_dim, last_dim, (1,), -1)
+        diag = (1 - (x * x).sum(-1)).sqrt().diag_embed()
+        x = x + diag
+        x = x.reshape((self.dim, self.dim))
+        return tril_matrix_to_vec(x @ x.T, -1)
+
+    def _compute_jacobian(self, x):
+        last_dim = self.dim * (self.dim - 1) // 2
+        if x.stop_gradient is not False:
+            x.stop_gradient = False
+        jacobian_matrix = []
+        outputs = self._tril_cholesky_to_tril_corr(x, last_dim)
+        for i in range(outputs.shape[0]):
+            grad = paddle.static.gradients(outputs[i], x)[0]
+            grad = grad.reshape((last_dim,))
+            jacobian_matrix.append(grad)
+        J = paddle.stack(jacobian_matrix, axis=0)
+        logabsdet = paddle.linalg.slogdet(J)
+        return logabsdet, J
+
+
+@parameterize.place(config.DEVICES)
+@parameterize.parameterize_cls(
+    (parameterize.TEST_CASE_NAME),
+    [
+        ('lkj_cholesky_test_err'),
+    ],
+)
+class LKJCholeskyTestError(unittest.TestCase):
+    def setUp(self):
+        self.program = paddle.static.Program()
+
+    @parameterize.parameterize_func(
+        [
+            (1, ValueError),  # dim < 2
+            (3.0, TypeError),  # dim is float
+        ]
+    )
+    def test_bad_parameter(self, dim, error):
+        with paddle.static.program_guard(self.program):
+            self.assertRaises(error, lkj_cholesky.LKJCholesky, dim)
+
+    @parameterize.parameterize_func([(10,)])  # not sequence object sample shape
+    def test_bad_sample_shape(self, shape):
+        with paddle.static.program_guard(self.program):
+            lkj = lkj_cholesky.LKJCholesky(3)
+            self.assertRaises(TypeError, lkj.sample, shape)
+
+
+if __name__ == '__main__':
+    unittest.main()

test/distribution/test_distribution_lognormal.py

@@ -0,0 +1,285 @@
+# Copyright (c) 2022 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
+#
+#     http://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 math
+import unittest
+
+import numpy as np
+import scipy.stats
+from distribution import config
+from parameterize import TEST_CASE_NAME, parameterize_cls, place, xrand
+from test_distribution import DistributionNumpy
+
+import paddle
+from paddle.distribution.kl import kl_divergence
+from paddle.distribution.lognormal import LogNormal
+from paddle.distribution.normal import Normal
+
+
+class LogNormalNumpy(DistributionNumpy):
+    def __init__(self, loc, scale):
+        self.loc = np.array(loc)
+        self.scale = np.array(scale)
+        if str(self.loc.dtype) not in ['float32', 'float64']:
+            self.loc = self.loc.astype('float32')
+            self.scale = self.scale.astype('float32')
+
+    @property
+    def mean(self):
+        var = self.scale * self.scale
+        return np.exp(self.loc + var / 2)
+
+    @property
+    def variance(self):
+        var = self.scale * self.scale
+        return (np.exp(var) - 1) * np.exp(2 * self.loc + var)
+
+    def log_prob(self, value):
+        var = self.scale * self.scale
+        log_scale = np.log(self.scale)
+        return (
+            -((np.log(value) - self.loc) * (np.log(value) - self.loc))
+            / (2.0 * var)
+            - log_scale
+            - math.log(math.sqrt(2.0 * math.pi))
+            - np.log(value)
+        )
+
+    def probs(self, value):
+        var = self.scale * self.scale
+        return np.exp(
+            -1.0
+            * ((np.log(value) - self.loc) * (np.log(value) - self.loc))
+            / (2.0 * var)
+        ) / (math.sqrt(2 * math.pi) * self.scale * value)
+
+    def entropy(self):
+        return (
+            0.5
+            + self.loc
+            + 0.5 * np.log(np.array(2.0 * math.pi).astype(self.loc.dtype))
+            + np.log(self.scale)
+        )
+
+    def kl_divergence(self, other):
+        var_ratio = self.scale / other.scale
+        var_ratio = var_ratio * var_ratio
+        t1 = (self.loc - other.loc) / other.scale
+        t1 = t1 * t1
+        return 0.5 * (var_ratio + t1 - 1 - np.log(var_ratio))
+
+
+@place(config.DEVICES)
+@parameterize_cls(
+    (TEST_CASE_NAME, 'loc', 'scale', 'value'),
+    [
+        ('one-dim', xrand((2,)), xrand((2,)), xrand((2,))),
+        ('multi-dim', xrand((3, 3)), xrand((3, 3)), xrand((3, 3))),
+    ],
+)
+class LogNormalTest(unittest.TestCase):
+    def setUp(self):
+        paddle.disable_static()
+        self.paddle_lognormal = LogNormal(
+            loc=paddle.to_tensor(self.loc), scale=paddle.to_tensor(self.scale)
+        )
+        self.np_lognormal = LogNormalNumpy(self.loc, self.scale)
+
+    def test_mean(self):
+        mean = self.paddle_lognormal.mean
+        np_mean = self.np_lognormal.mean
+        self.assertEqual(mean.numpy().dtype, np_mean.dtype)
+        np.testing.assert_allclose(
+            mean,
+            np_mean,
+            rtol=config.RTOL.get(str(self.scale.dtype)),
+            atol=config.ATOL.get(str(self.scale.dtype)),
+        )
+
+    def test_variance(self):
+        var = self.paddle_lognormal.variance
+        np_var = self.np_lognormal.variance
+        self.assertEqual(var.numpy().dtype, np_var.dtype)
+        np.testing.assert_allclose(
+            var,
+            np_var,
+            rtol=config.RTOL.get(str(self.scale.dtype)),
+            atol=config.ATOL.get(str(self.scale.dtype)),
+        )
+
+    def test_entropy(self):
+        entropy = self.paddle_lognormal.entropy()
+        np_entropy = self.np_lognormal.entropy()
+        self.assertEqual(entropy.numpy().dtype, np_entropy.dtype)
+        np.testing.assert_allclose(
+            entropy,
+            np_entropy,
+            rtol=config.RTOL.get(str(self.scale.dtype)),
+            atol=config.ATOL.get(str(self.scale.dtype)),
+        )
+
+    def test_probs(self):
+        with paddle.base.dygraph.guard(self.place):
+            probs = self.paddle_lognormal.probs(paddle.to_tensor(self.value))
+            np_probs = self.np_lognormal.probs(self.value)
+            np.testing.assert_allclose(
+                probs,
+                np_probs,
+                rtol=config.RTOL.get(str(self.scale.dtype)),
+                atol=config.ATOL.get(str(self.scale.dtype)),
+            )
+
+    def test_log_prob(self):
+        with paddle.base.dygraph.guard(self.place):
+            log_prob = self.paddle_lognormal.log_prob(
+                paddle.to_tensor(self.value)
+            )
+            np_log_prob = self.np_lognormal.log_prob(self.value)
+            np.testing.assert_allclose(
+                log_prob,
+                np_log_prob,
+                rtol=config.RTOL.get(str(self.scale.dtype)),
+                atol=config.ATOL.get(str(self.scale.dtype)),
+            )
+
+
+@place(config.DEVICES)
+@parameterize_cls(
+    (TEST_CASE_NAME, 'loc', 'scale'),
+    [('sample', xrand((4,)), xrand((4,), min=0, max=1))],
+)
+class TestLogNormalSample(unittest.TestCase):
+    def setUp(self):
+        paddle.disable_static()
+        self.paddle_lognormal = LogNormal(loc=self.loc, scale=self.scale)
+        n = 1000000
+        self.sample_shape = (n,)
+        self.rsample_shape = (n,)
+        self.samples = self.paddle_lognormal.sample(self.sample_shape)
+        self.rsamples = self.paddle_lognormal.rsample(self.rsample_shape)
+
+    def test_sample(self):
+        samples_mean = self.samples.mean(axis=0)
+        samples_var = self.samples.var(axis=0)
+        np.testing.assert_allclose(
+            samples_mean, self.paddle_lognormal.mean, rtol=0.1, atol=0
+        )
+        np.testing.assert_allclose(
+            samples_var, self.paddle_lognormal.variance, rtol=0.1, atol=0
+        )
+
+        rsamples_mean = self.rsamples.mean(axis=0)
+        rsamples_var = self.rsamples.var(axis=0)
+        np.testing.assert_allclose(
+            rsamples_mean, self.paddle_lognormal.mean, rtol=0.1, atol=0
+        )
+        np.testing.assert_allclose(
+            rsamples_var, self.paddle_lognormal.variance, rtol=0.1, atol=0
+        )
+
+        batch_shape = (self.loc + self.scale).shape
+        self.assertEqual(
+            self.samples.shape, list(self.sample_shape + batch_shape)
+        )
+        self.assertEqual(
+            self.rsamples.shape, list(self.rsample_shape + batch_shape)
+        )
+
+        for i in range(len(self.scale)):
+            self.assertTrue(
+                self._kstest(self.loc[i], self.scale[i], self.samples[:, i])
+            )
+            self.assertTrue(
+                self._kstest(self.loc[i], self.scale[i], self.rsamples[:, i])
+            )
+
+    def _kstest(self, loc, scale, samples):
+        # Uses the Kolmogorov-Smirnov test for goodness of fit.
+        ks, _ = scipy.stats.kstest(
+            samples, scipy.stats.lognorm(s=scale, scale=np.exp(loc)).cdf
+        )
+        return ks < 0.02
+
+
+@place(config.DEVICES)
+@parameterize_cls(
+    (TEST_CASE_NAME, 'loc1', 'scale1', 'loc2', 'scale2'),
+    [
+        ('one-dim', xrand((2,)), xrand((2,)), xrand((2,)), xrand((2,))),
+        (
+            'multi-dim',
+            xrand((2, 2)),
+            xrand((2, 2)),
+            xrand((2, 2)),
+            xrand((2, 2)),
+        ),
+    ],
+)
+class TestLogNormalKL(unittest.TestCase):
+    def setUp(self):
+        paddle.disable_static()
+        self.ln_a = LogNormal(
+            loc=paddle.to_tensor(self.loc1), scale=paddle.to_tensor(self.scale1)
+        )
+        self.ln_b = LogNormal(
+            loc=paddle.to_tensor(self.loc2), scale=paddle.to_tensor(self.scale2)
+        )
+        self.normal_a = Normal(
+            loc=paddle.to_tensor(self.loc1), scale=paddle.to_tensor(self.scale1)
+        )
+        self.normal_b = Normal(
+            loc=paddle.to_tensor(self.loc2), scale=paddle.to_tensor(self.scale2)
+        )
+
+    def test_kl_divergence(self):
+        kl0 = self.ln_a.kl_divergence(self.ln_b)
+        kl1 = kl_divergence(self.ln_a, self.ln_b)
+        kl_normal = kl_divergence(self.normal_a, self.normal_b)
+        kl_formula = self._kl(self.ln_a, self.ln_b)
+
+        self.assertEqual(tuple(kl0.shape), self.scale1.shape)
+        self.assertEqual(tuple(kl1.shape), self.scale1.shape)
+        np.testing.assert_allclose(
+            kl0,
+            kl_formula,
+            rtol=config.RTOL.get(str(self.scale1.dtype)),
+            atol=config.ATOL.get(str(self.scale1.dtype)),
+        )
+        np.testing.assert_allclose(
+            kl1,
+            kl_formula,
+            rtol=config.RTOL.get(str(self.scale1.dtype)),
+            atol=config.ATOL.get(str(self.scale1.dtype)),
+        )
+        np.testing.assert_allclose(
+            kl_normal,
+            kl_formula,
+            rtol=config.RTOL.get(str(self.scale1.dtype)),
+            atol=config.ATOL.get(str(self.scale1.dtype)),
+        )
+
+    def _kl(self, dist1, dist2):
+        loc1 = np.array(dist1.loc)
+        loc2 = np.array(dist2.loc)
+        scale1 = np.array(dist1.scale)
+        scale2 = np.array(dist2.scale)
+        var_ratio = scale1 / scale2
+        var_ratio = var_ratio * var_ratio
+        t1 = (loc1 - loc2) / scale2
+        t1 = t1 * t1
+        return 0.5 * (var_ratio + t1 - 1 - np.log(var_ratio))
+
+
+if __name__ == '__main__':
+    unittest.main()

test/distribution/test_distribution_lognormal_static.py

@@ -0,0 +1,298 @@
+# Copyright (c) 2022 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
+#
+#     http://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
+import scipy.stats
+from distribution import config
+from parameterize import TEST_CASE_NAME, parameterize_cls, place, xrand
+from test_distribution_lognormal import LogNormalNumpy
+
+import paddle
+from paddle.distribution.kl import kl_divergence
+from paddle.distribution.lognormal import LogNormal
+from paddle.distribution.normal import Normal
+
+
+@place(config.DEVICES)
+@parameterize_cls(
+    (TEST_CASE_NAME, 'loc', 'scale', 'value'),
+    [
+        ('one-dim', xrand((2,)), xrand((2,)), xrand((2,))),
+        ('multi-dim', xrand((3, 3)), xrand((3, 3)), xrand((3, 3))),
+    ],
+    test_pir=True,
+)
+class TestLogNormal(unittest.TestCase):
+    def run_program(self):
+        paddle.enable_static()
+        startup_program = paddle.static.Program()
+        main_program = paddle.static.Program()
+        executor = paddle.static.Executor(self.place)
+        with paddle.static.program_guard(main_program, startup_program):
+            loc = paddle.static.data('loc', self.loc.shape, self.loc.dtype)
+            scale = paddle.static.data(
+                'scale', self.scale.shape, self.scale.dtype
+            )
+            value = paddle.static.data(
+                'value', self.value.shape, self.value.dtype
+            )
+            self.paddle_lognormal = LogNormal(loc=loc, scale=scale)
+            self.np_lognormal = LogNormalNumpy(loc=self.loc, scale=self.scale)
+            mean = self.paddle_lognormal.mean
+            var = self.paddle_lognormal.variance
+            entropy = self.paddle_lognormal.entropy()
+            probs = self.paddle_lognormal.probs(value)
+            log_prob = self.paddle_lognormal.log_prob(value)
+        fetch_list = [mean, var, entropy, probs, log_prob]
+        self.feeds = {'loc': self.loc, 'scale': self.scale, 'value': self.value}
+
+        executor.run(startup_program)
+        [
+            self.mean,
+            self.var,
+            self.entropy,
+            self.probs,
+            self.log_prob,
+        ] = executor.run(main_program, feed=self.feeds, fetch_list=fetch_list)
+
+    def setUp(self):
+        if self.test_pir:
+            with paddle.pir_utils.IrGuard():
+                self.run_program()
+        else:
+            self.run_program()
+
+    def test_mean(self):
+        np_mean = self.np_lognormal.mean
+        self.assertEqual(str(self.mean.dtype).split('.')[-1], self.scale.dtype)
+        np.testing.assert_allclose(
+            self.mean,
+            np_mean,
+            rtol=config.RTOL.get(str(self.scale.dtype)),
+            atol=config.ATOL.get(str(self.scale.dtype)),
+        )
+
+    def test_var(self):
+        np_var = self.np_lognormal.variance
+        self.assertEqual(str(self.var.dtype).split('.')[-1], self.scale.dtype)
+        np.testing.assert_allclose(
+            self.var,
+            np_var,
+            rtol=config.RTOL.get(str(self.scale.dtype)),
+            atol=config.ATOL.get(str(self.scale.dtype)),
+        )
+
+    def test_entropy(self):
+        np_entropy = self.np_lognormal.entropy()
+        self.assertEqual(
+            str(self.entropy.dtype).split('.')[-1], self.scale.dtype
+        )
+        np.testing.assert_allclose(
+            self.entropy,
+            np_entropy,
+            rtol=config.RTOL.get(str(self.scale.dtype)),
+            atol=config.ATOL.get(str(self.scale.dtype)),
+        )
+
+    def test_probs(self):
+        np_probs = self.np_lognormal.probs(self.value)
+        np.testing.assert_allclose(
+            self.probs,
+            np_probs,
+            rtol=config.RTOL.get(str(self.scale.dtype)),
+            atol=config.ATOL.get(str(self.scale.dtype)),
+        )
+
+    def test_log_prob(self):
+        np_log_prob = self.np_lognormal.log_prob(self.value)
+        np.testing.assert_allclose(
+            self.log_prob,
+            np_log_prob,
+            rtol=config.RTOL.get(str(self.scale.dtype)),
+            atol=config.ATOL.get(str(self.scale.dtype)),
+        )
+
+
+@place(config.DEVICES)
+@parameterize_cls(
+    (TEST_CASE_NAME, 'loc', 'scale'),
+    [('sample', xrand((4,)), xrand((4,), min=0, max=1))],
+    test_pir=True,
+)
+class TestLogNormalSample(unittest.TestCase):
+    def run_program(self):
+        paddle.enable_static()
+        startup_program = paddle.static.Program()
+        main_program = paddle.static.Program()
+        executor = paddle.static.Executor(self.place)
+        with paddle.static.program_guard(main_program, startup_program):
+            loc = paddle.static.data('loc', self.loc.shape, self.loc.dtype)
+            scale = paddle.static.data(
+                'scale', self.scale.shape, self.scale.dtype
+            )
+            n = 1000000
+            self.sample_shape = (n,)
+            self.rsample_shape = (n,)
+            self.paddle_lognormal = LogNormal(loc=loc, scale=scale)
+            mean = self.paddle_lognormal.mean
+            variance = self.paddle_lognormal.variance
+            samples = self.paddle_lognormal.sample(self.sample_shape)
+            rsamples = self.paddle_lognormal.rsample(self.rsample_shape)
+        fetch_list = [mean, variance, samples, rsamples]
+        self.feeds = {'loc': self.loc, 'scale': self.scale}
+
+        executor.run(startup_program)
+        [self.mean, self.variance, self.samples, self.rsamples] = executor.run(
+            main_program, feed=self.feeds, fetch_list=fetch_list
+        )
+
+    def setUp(self):
+        if self.test_pir:
+            with paddle.pir_utils.IrGuard():
+                self.run_program()
+        else:
+            self.run_program()
+
+    def test_sample(self):
+        samples_mean = self.samples.mean(axis=0)
+        samples_var = self.samples.var(axis=0)
+        np.testing.assert_allclose(samples_mean, self.mean, rtol=0.1, atol=0)
+        np.testing.assert_allclose(samples_var, self.variance, rtol=0.1, atol=0)
+
+        rsamples_mean = self.rsamples.mean(axis=0)
+        rsamples_var = self.rsamples.var(axis=0)
+        np.testing.assert_allclose(rsamples_mean, self.mean, rtol=0.1, atol=0)
+        np.testing.assert_allclose(
+            rsamples_var, self.variance, rtol=0.1, atol=0
+        )
+
+        batch_shape = (self.loc + self.scale).shape
+        self.assertEqual(self.samples.shape, self.sample_shape + batch_shape)
+        self.assertEqual(self.rsamples.shape, self.rsample_shape + batch_shape)
+
+        for i in range(len(self.scale)):
+            self.assertTrue(
+                self._kstest(self.loc[i], self.scale[i], self.samples[:, i])
+            )
+            self.assertTrue(
+                self._kstest(self.loc[i], self.scale[i], self.rsamples[:, i])
+            )
+
+    def _kstest(self, loc, scale, samples):
+        # Uses the Kolmogorov-Smirnov test for goodness of fit.
+        ks, _ = scipy.stats.kstest(
+            samples, scipy.stats.lognorm(s=scale, scale=np.exp(loc)).cdf
+        )
+        return ks < 0.02
+
+
+@place(config.DEVICES)
+@parameterize_cls(
+    (TEST_CASE_NAME, 'loc1', 'scale1', 'loc2', 'scale2'),
+    [
+        ('one-dim', xrand((2,)), xrand((2,)), xrand((2,)), xrand((2,))),
+        (
+            'multi-dim',
+            xrand((2, 2)),
+            xrand((2, 2)),
+            xrand((2, 2)),
+            xrand((2, 2)),
+        ),
+    ],
+    test_pir=True,
+)
+class TestLogNormalKL(unittest.TestCase):
+    def run_program(self):
+        paddle.enable_static()
+        startup_program = paddle.static.Program()
+        main_program = paddle.static.Program()
+        executor = paddle.static.Executor(self.place)
+        with paddle.static.program_guard(main_program, startup_program):
+            loc1 = paddle.static.data('loc1', self.loc1.shape, self.loc1.dtype)
+            scale1 = paddle.static.data(
+                'scale1', self.scale1.shape, self.scale1.dtype
+            )
+            loc2 = paddle.static.data('loc2', self.loc2.shape, self.loc2.dtype)
+            scale2 = paddle.static.data(
+                'scale2', self.scale2.shape, self.scale2.dtype
+            )
+
+            self.ln_a = LogNormal(loc=loc1, scale=scale1)
+            self.ln_b = LogNormal(loc=loc2, scale=scale2)
+            self.normal_a = Normal(loc=loc1, scale=scale1)
+            self.normal_b = Normal(loc=loc2, scale=scale2)
+
+            kl0 = self.ln_a.kl_divergence(self.ln_b)
+            kl1 = kl_divergence(self.ln_a, self.ln_b)
+            kl_normal = kl_divergence(self.normal_a, self.normal_b)
+            kl_formula = self._kl(self.ln_a, self.ln_b)
+
+        fetch_list = [kl0, kl1, kl_normal, kl_formula]
+        self.feeds = {
+            'loc1': self.loc1,
+            'scale1': self.scale1,
+            'loc2': self.loc2,
+            'scale2': self.scale2,
+        }
+
+        executor.run(startup_program)
+        [self.kl0, self.kl1, self.kl_normal, self.kl_formula] = executor.run(
+            main_program, feed=self.feeds, fetch_list=fetch_list
+        )
+
+    def setUp(self):
+        if self.test_pir:
+            with paddle.pir_utils.IrGuard():
+                self.run_program()
+        else:
+            self.run_program()
+
+    def test_kl_divergence(self):
+        np.testing.assert_allclose(
+            self.kl0,
+            self.kl_formula,
+            rtol=config.RTOL.get(str(self.scale1.dtype)),
+            atol=config.ATOL.get(str(self.scale1.dtype)),
+        )
+
+        np.testing.assert_allclose(
+            self.kl1,
+            self.kl_formula,
+            rtol=config.RTOL.get(str(self.scale1.dtype)),
+            atol=config.ATOL.get(str(self.scale1.dtype)),
+        )
+
+        np.testing.assert_allclose(
+            self.kl_normal,
+            self.kl_formula,
+            rtol=config.RTOL.get(str(self.scale1.dtype)),
+            atol=config.ATOL.get(str(self.scale1.dtype)),
+        )
+
+    def _kl(self, dist1, dist2):
+        loc1 = dist1.loc
+        loc2 = dist2.loc
+        scale1 = dist1.scale
+        scale2 = dist2.scale
+        var_ratio = scale1 / scale2
+        var_ratio = var_ratio * var_ratio
+        t1 = (loc1 - loc2) / scale2
+        t1 = t1 * t1
+        return 0.5 * (var_ratio + t1 - 1 - np.log(var_ratio))
+
+
+if __name__ == '__main__':
+    unittest.main()

test/distribution/test_distribution_multinomial.py

@@ -0,0 +1,168 @@
+# Copyright (c) 2021 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
+#
+#     http://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
+import parameterize
+import scipy.stats
+from distribution import config
+
+import paddle
+
+
+@parameterize.place(config.DEVICES)
+@parameterize.parameterize_cls(
+    (parameterize.TEST_CASE_NAME, 'total_count', 'probs'),
+    [
+        ('one-dim', 10, parameterize.xrand((3,))),
+        ('multi-dim', 9, parameterize.xrand((10, 20))),
+        ('prob-sum-one', 10, np.array([0.5, 0.2, 0.3])),
+        ('prob-sum-non-one', 10, np.array([2.0, 3.0, 5.0])),
+    ],
+)
+class TestMultinomial(unittest.TestCase):
+    def setUp(self):
+        self._dist = paddle.distribution.Multinomial(
+            total_count=self.total_count, probs=paddle.to_tensor(self.probs)
+        )
+
+    def test_mean(self):
+        mean = self._dist.mean
+        self.assertEqual(mean.numpy().dtype, self.probs.dtype)
+        np.testing.assert_allclose(
+            mean,
+            self._np_mean(),
+            rtol=config.RTOL.get(str(self.probs.dtype)),
+            atol=config.ATOL.get(str(self.probs.dtype)),
+        )
+
+    def test_variance(self):
+        var = self._dist.variance
+        self.assertEqual(var.numpy().dtype, self.probs.dtype)
+        np.testing.assert_allclose(
+            var,
+            self._np_variance(),
+            rtol=config.RTOL.get(str(self.probs.dtype)),
+            atol=config.ATOL.get(str(self.probs.dtype)),
+        )
+
+    def test_entropy(self):
+        entropy = self._dist.entropy()
+        self.assertEqual(entropy.numpy().dtype, self.probs.dtype)
+        np.testing.assert_allclose(
+            entropy,
+            self._np_entropy(),
+            rtol=config.RTOL.get(str(self.probs.dtype)),
+            atol=config.ATOL.get(str(self.probs.dtype)),
+        )
+
+    def test_sample(self):
+        sample_shape = ()
+        samples = self._dist.sample(sample_shape)
+        self.assertEqual(samples.numpy().dtype, self.probs.dtype)
+        self.assertEqual(
+            tuple(samples.shape),
+            sample_shape + self._dist.batch_shape + self._dist.event_shape,
+        )
+
+        sample_shape = (6,)
+        samples = self._dist.sample(sample_shape)
+        self.assertEqual(samples.numpy().dtype, self.probs.dtype)
+        self.assertEqual(
+            tuple(samples.shape),
+            sample_shape + self._dist.batch_shape + self._dist.event_shape,
+        )
+        self.assertTrue(
+            np.all(samples.sum(-1).numpy() == self._dist.total_count)
+        )
+
+        sample_shape = (5000,)
+        samples = self._dist.sample(sample_shape)
+        sample_mean = samples.mean(axis=0)
+        # Tolerance value 0.2 is empirical value which is consistent with
+        # TensorFlow
+        np.testing.assert_allclose(
+            sample_mean, self._dist.mean, atol=0, rtol=0.20
+        )
+
+    def _np_variance(self):
+        probs = self.probs / self.probs.sum(-1, keepdims=True)
+        return self.total_count * probs * (1 - probs)
+
+    def _np_mean(self):
+        probs = self.probs / self.probs.sum(-1, keepdims=True)
+        return self.total_count * probs
+
+    def _np_entropy(self):
+        probs = self.probs / self.probs.sum(-1, keepdims=True)
+        return scipy.stats.multinomial.entropy(self.total_count, probs)
+
+
+@parameterize.place(config.DEVICES)
+@parameterize.parameterize_cls(
+    (parameterize.TEST_CASE_NAME, 'total_count', 'probs', 'value'),
+    [
+        (
+            'value-float',
+            10,
+            np.array([0.2, 0.3, 0.5]),
+            np.array([2.0, 3.0, 5.0]),
+        ),
+        ('value-int', 10, np.array([0.2, 0.3, 0.5]), np.array([2, 3, 5])),
+        (
+            'value-multi-dim',
+            10,
+            np.array([[0.3, 0.7], [0.5, 0.5]]),
+            np.array([[4.0, 6], [8, 2]]),
+        ),
+        # ('value-sum-non-n', 10, np.array([0.5, 0.2, 0.3]), np.array([4,5,2])),
+    ],
+)
+class TestMultinomialPmf(unittest.TestCase):
+    def setUp(self):
+        self._dist = paddle.distribution.Multinomial(
+            total_count=self.total_count, probs=paddle.to_tensor(self.probs)
+        )
+
+    def test_prob(self):
+        np.testing.assert_allclose(
+            self._dist.prob(paddle.to_tensor(self.value)),
+            scipy.stats.multinomial.pmf(
+                self.value, self.total_count, self.probs
+            ),
+            rtol=config.RTOL.get(str(self.probs.dtype)),
+            atol=config.ATOL.get(str(self.probs.dtype)),
+        )
+
+
+@parameterize.place(config.DEVICES)
+@parameterize.parameterize_cls(
+    (config.TEST_CASE_NAME, 'total_count', 'probs'),
+    [
+        ('total_count_le_one', 0, np.array([0.3, 0.7])),
+        ('total_count_float', np.array([0.3, 0.7])),
+        ('probs_zero_dim', np.array(0)),
+    ],
+)
+class TestMultinomialException(unittest.TestCase):
+    def TestInit(self):
+        with self.assertRaises(ValueError):
+            paddle.distribution.Multinomial(
+                self.total_count, paddle.to_tensor(self.probs)
+            )
+
+
+if __name__ == '__main__':
+    unittest.main()

test/distribution/test_distribution_multinomial_static.py

@@ -0,0 +1,199 @@
+# Copyright (c) 2021 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
+#
+#     http://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
+import parameterize
+import scipy.stats
+from distribution import config
+
+import paddle
+
+paddle.enable_static()
+
+
+@parameterize.place(config.DEVICES)
+@parameterize.parameterize_cls(
+    (parameterize.TEST_CASE_NAME, 'total_count', 'probs'),
+    [
+        ('one-dim', 5, parameterize.xrand((3,))),
+        ('multi-dim', 9, parameterize.xrand((2, 3))),
+        ('prob-sum-one', 5, np.array([0.5, 0.2, 0.3])),
+        ('prob-sum-non-one', 5, np.array([2.0, 3.0, 5.0])),
+    ],
+)
+class TestMultinomial(unittest.TestCase):
+    def setUp(self):
+        startup_program = paddle.static.Program()
+        main_program = paddle.static.Program()
+        executor = paddle.static.Executor(self.place)
+        with paddle.static.program_guard(main_program, startup_program):
+            probs = paddle.static.data(
+                'probs', self.probs.shape, self.probs.dtype
+            )
+            dist = paddle.distribution.Multinomial(self.total_count, probs)
+            mean = dist.mean
+            var = dist.variance
+            entropy = dist.entropy()
+            mini_samples = dist.sample(shape=(6,))
+            large_samples = dist.sample(shape=(5000,))
+        fetch_list = [mean, var, entropy, mini_samples, large_samples]
+        feed = {'probs': self.probs}
+
+        executor.run(startup_program)
+        [
+            self.mean,
+            self.var,
+            self.entropy,
+            self.mini_samples,
+            self.large_samples,
+        ] = executor.run(main_program, feed=feed, fetch_list=fetch_list)
+
+    def test_mean(self):
+        self.assertEqual(str(self.mean.dtype).split('.')[-1], self.probs.dtype)
+        np.testing.assert_allclose(
+            self.mean,
+            self._np_mean(),
+            rtol=config.RTOL.get(str(self.probs.dtype)),
+            atol=config.ATOL.get(str(self.probs.dtype)),
+        )
+
+    def test_variance(self):
+        self.assertEqual(str(self.var.dtype).split('.')[-1], self.probs.dtype)
+        np.testing.assert_allclose(
+            self.var,
+            self._np_variance(),
+            rtol=config.RTOL.get(str(self.probs.dtype)),
+            atol=config.ATOL.get(str(self.probs.dtype)),
+        )
+
+    def test_entropy(self):
+        self.assertEqual(
+            str(self.entropy.dtype).split('.')[-1], self.probs.dtype
+        )
+        np.testing.assert_allclose(
+            self.entropy,
+            self._np_entropy(),
+            rtol=config.RTOL.get(str(self.probs.dtype)),
+            atol=config.ATOL.get(str(self.probs.dtype)),
+        )
+
+    def test_sample(self):
+        self.assertEqual(
+            str(self.mini_samples.dtype).split('.')[-1], self.probs.dtype
+        )
+        self.assertTrue(np.all(self.mini_samples.sum(-1) == self.total_count))
+
+        sample_mean = self.large_samples.mean(axis=0)
+        np.testing.assert_allclose(sample_mean, self.mean, atol=0, rtol=0.20)
+
+    def _np_variance(self):
+        probs = self.probs / self.probs.sum(-1, keepdims=True)
+        return self.total_count * probs * (1 - probs)
+
+    def _np_mean(self):
+        probs = self.probs / self.probs.sum(-1, keepdims=True)
+        return self.total_count * probs
+
+    def _np_entropy(self):
+        probs = self.probs / self.probs.sum(-1, keepdims=True)
+        return scipy.stats.multinomial.entropy(self.total_count, probs)
+
+
+@parameterize.place(config.DEVICES)
+@parameterize.parameterize_cls(
+    (parameterize.TEST_CASE_NAME, 'total_count', 'probs', 'value'),
+    [
+        (
+            'value-float',
+            5,
+            np.array([0.2, 0.3, 0.5]),
+            np.array([1.0, 1.0, 3.0]),
+        ),
+        ('value-int', 5, np.array([0.2, 0.3, 0.5]), np.array([2, 2, 1])),
+        (
+            'value-multi-dim',
+            5,
+            np.array([[0.3, 0.7], [0.5, 0.5]]),
+            np.array([[1.0, 4.0], [2.0, 3.0]]),
+        ),
+        # ('value-sum-non-n', 10, np.array([0.5, 0.2, 0.3]), np.array([4,5,2])),
+    ],
+)
+class TestMultinomialPmf(unittest.TestCase):
+    def setUp(self):
+        startup_program = paddle.static.Program()
+        main_program = paddle.static.Program()
+        executor = paddle.static.Executor(self.place)
+
+        with paddle.static.program_guard(main_program, startup_program):
+            probs = paddle.static.data(
+                'probs', self.probs.shape, self.probs.dtype
+            )
+            value = paddle.static.data(
+                'value', self.value.shape, self.value.dtype
+            )
+            dist = paddle.distribution.Multinomial(self.total_count, probs)
+            pmf = dist.prob(value)
+        feed = {'probs': self.probs, 'value': self.value}
+        fetch_list = [pmf]
+
+        executor.run(startup_program)
+        [self.pmf] = executor.run(
+            main_program, feed=feed, fetch_list=fetch_list
+        )
+
+    def test_prob(self):
+        np.testing.assert_allclose(
+            self.pmf,
+            scipy.stats.multinomial.pmf(
+                self.value, self.total_count, self.probs
+            ),
+            rtol=config.RTOL.get(str(self.probs.dtype)),
+            atol=config.ATOL.get(str(self.probs.dtype)),
+        )
+
+
+@parameterize.place(config.DEVICES)
+@parameterize.parameterize_cls(
+    (parameterize.TEST_CASE_NAME, 'total_count', 'probs'),
+    [
+        ('total_count_le_one', 0, np.array([0.3, 0.7])),
+        ('total_count_float', np.array([0.3, 0.7])),
+        ('probs_zero_dim', np.array(0)),
+    ],
+)
+class TestMultinomialException(unittest.TestCase):
+    def setUp(self):
+        startup_program = paddle.static.Program()
+        self.main_program = paddle.static.Program()
+        self.executor = paddle.static.Executor(self.place)
+
+        with paddle.static.program_guard(self.main_program, startup_program):
+            probs = paddle.static.data(
+                'probs', self.probs.shape, self.probs.dtype
+            )
+            dist = paddle.distribution.Multinomial(self.total_count, probs)
+        self.feed = {'probs': self.probs}
+
+        self.executor.run(startup_program)
+
+    def TestInit(self):
+        with self.assertRaises(ValueError):
+            self.executor.run(self.main_program, feed=self.feed, fetch=[])
+
+
+if __name__ == '__main__':
+    unittest.main()

test/distribution/test_distribution_multivariate_normal.py

@@ -0,0 +1,465 @@
+# Copyright (c) 2021 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
+#
+#     http://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
+import parameterize
+import scipy
+from distribution import config
+from parameterize import (
+    TEST_CASE_NAME,
+    parameterize_cls,
+)
+
+import paddle
+from paddle.distribution import constraint
+from paddle.distribution.multivariate_normal import MultivariateNormal
+
+
+@parameterize.place(config.DEVICES)
+@parameterize.parameterize_cls(
+    (parameterize.TEST_CASE_NAME, 'loc', 'covariance_matrix'),
+    [
+        (
+            'one-batch',
+            parameterize.xrand((2,), dtype='float32', min=1, max=2),
+            np.array([[2.0, 1.0], [1.0, 2.0]]),
+        ),
+        (
+            'multi-batch',
+            parameterize.xrand((2, 3), dtype='float64', min=-2, max=-1),
+            np.array([[4.0, 2.5, 2.0], [2.5, 3.0, 1.2], [2.0, 1.2, 4.0]]),
+        ),
+    ],
+)
+class TestMVN(unittest.TestCase):
+    def setUp(self):
+        self._dist = MultivariateNormal(
+            loc=paddle.to_tensor(self.loc),
+            covariance_matrix=paddle.to_tensor(self.covariance_matrix),
+        )
+
+    def test_mean(self):
+        mean = self._dist.mean
+        self.assertEqual(mean.numpy().dtype, self.loc.dtype)
+        np.testing.assert_allclose(
+            mean,
+            self._np_mean(),
+            rtol=config.RTOL.get(str(self.loc.dtype)),
+            atol=config.ATOL.get(str(self.loc.dtype)),
+        )
+
+    def test_variance(self):
+        var = self._dist.variance
+        self.assertEqual(var.numpy().dtype, self.loc.dtype)
+        np.testing.assert_allclose(
+            var,
+            self._np_variance(),
+            rtol=config.RTOL.get(str(self.loc.dtype)),
+            atol=config.ATOL.get(str(self.loc.dtype)),
+        )
+
+    def test_entropy(self):
+        entropy = self._dist.entropy()
+        self.assertEqual(entropy.numpy().dtype, self.loc.dtype)
+        np.testing.assert_allclose(
+            entropy,
+            self._np_entropy(),
+            rtol=config.RTOL.get(str(self.loc.dtype)),
+            atol=config.ATOL.get(str(self.loc.dtype)),
+        )
+
+    def test_sample(self):
+        sample_shape = ()
+        samples = self._dist.sample(sample_shape)
+        self.assertEqual(samples.numpy().dtype, self.loc.dtype)
+        self.assertEqual(
+            tuple(samples.shape),
+            sample_shape + self._dist.batch_shape + self._dist.event_shape,
+        )
+
+        sample_shape = (50000,)
+        samples = self._dist.sample(sample_shape)
+        sample_mean = samples.mean(axis=0)
+        sample_variance = samples.var(axis=0)
+
+        # `atol` and `rtol` refer to ``test_distribution_normal`` and ``test_distribution_lognormal``
+        np.testing.assert_allclose(
+            sample_mean, self._dist.mean, atol=0.0, rtol=0.1
+        )
+        np.testing.assert_allclose(
+            sample_variance, self._dist.variance, atol=0.0, rtol=0.1
+        )
+
+    def _np_variance(self):
+        batch_shape = np.broadcast_shapes(
+            self.covariance_matrix.shape[:-2], self.loc.shape[:-1]
+        )
+        event_shape = self.loc.shape[-1:]
+        return np.broadcast_to(
+            np.diag(self.covariance_matrix), batch_shape + event_shape
+        )
+
+    def _np_mean(self):
+        return self.loc
+
+    def _np_entropy(self):
+        if len(self.loc.shape) <= 1:
+            return scipy.stats.multivariate_normal.entropy(
+                self.loc, self.covariance_matrix
+            )
+        else:
+            return np.apply_along_axis(
+                lambda i: scipy.stats.multivariate_normal.entropy(
+                    i, self.covariance_matrix
+                ),
+                axis=1,
+                arr=self.loc,
+            )
+
+
+@parameterize.place(config.DEVICES)
+@parameterize.parameterize_cls(
+    (parameterize.TEST_CASE_NAME, 'loc', 'precision_matrix', 'value'),
+    [
+        (
+            'value-same-shape',
+            parameterize.xrand((2,), dtype='float32', min=-2, max=2),
+            np.array([[2.0, 1.0], [1.0, 2.0]]),
+            parameterize.xrand((2,), dtype='float32', min=-5, max=5),
+        ),
+        (
+            'value-broadcast-shape',
+            parameterize.xrand((2,), dtype='float64', min=-2, max=2),
+            np.array([[2.0, 1.0], [1.0, 2.0]]),
+            parameterize.xrand((3, 2), dtype='float64', min=-5, max=5),
+        ),
+    ],
+)
+class TestMVNProbs(unittest.TestCase):
+    def setUp(self):
+        self._dist = MultivariateNormal(
+            loc=paddle.to_tensor(self.loc),
+            precision_matrix=paddle.to_tensor(self.precision_matrix),
+        )
+        self.cov = np.linalg.inv(self.precision_matrix)
+
+    def test_prob(self):
+        if len(self.value.shape) <= 1:
+            scipy_pdf = scipy.stats.multivariate_normal.pdf(
+                self.value, self.loc, self.cov
+            )
+        else:
+            scipy_pdf = np.apply_along_axis(
+                lambda i: scipy.stats.multivariate_normal.pdf(
+                    i, self.loc, self.cov
+                ),
+                axis=1,
+                arr=self.value,
+            )
+        np.testing.assert_allclose(
+            self._dist.prob(paddle.to_tensor(self.value)),
+            scipy_pdf,
+            rtol=config.RTOL.get(str(self.loc.dtype)),
+            atol=config.ATOL.get(str(self.loc.dtype)),
+        )
+
+    def test_log_prob(self):
+        if len(self.value.shape) <= 1:
+            scipy_logpdf = scipy.stats.multivariate_normal.logpdf(
+                self.value, self.loc, self.cov
+            )
+        else:
+            scipy_logpdf = np.apply_along_axis(
+                lambda i: scipy.stats.multivariate_normal.logpdf(
+                    i, self.loc, self.cov
+                ),
+                axis=1,
+                arr=self.value,
+            )
+        np.testing.assert_allclose(
+            self._dist.log_prob(paddle.to_tensor(self.value)),
+            scipy_logpdf,
+            rtol=config.RTOL.get(str(self.loc.dtype)),
+            atol=config.ATOL.get(str(self.loc.dtype)),
+        )
+
+
+@parameterize.place(config.DEVICES)
+@parameterize.parameterize_cls(
+    (parameterize.TEST_CASE_NAME, 'mu_1', 'tril_1', 'mu_2', 'tril_2'),
+    [
+        (
+            'one-batch',
+            parameterize.xrand((2,), dtype='float32', min=-2, max=2),
+            np.array([[2.0, 0.0], [1.0, 2.0]]),
+            parameterize.xrand((2,), dtype='float32', min=-2, max=2),
+            np.array([[3.0, 0.0], [2.0, 3.0]]),
+        )
+    ],
+)
+class TestMVNKL(unittest.TestCase):
+    def setUp(self):
+        paddle.disable_static()
+        self._dist1 = MultivariateNormal(
+            loc=paddle.to_tensor(self.mu_1),
+            scale_tril=paddle.to_tensor(self.tril_1),
+        )
+        self._dist2 = MultivariateNormal(
+            loc=paddle.to_tensor(self.mu_2),
+            scale_tril=paddle.to_tensor(self.tril_2),
+        )
+
+    def test_kl_divergence(self):
+        kl0 = self._dist1.kl_divergence(self._dist2)
+        kl1 = self.kl_divergence(self._dist1, self._dist2)
+
+        self.assertEqual(tuple(kl0.shape), self._dist1.batch_shape)
+        self.assertEqual(tuple(kl1.shape), self._dist1.batch_shape)
+        np.testing.assert_allclose(
+            kl0,
+            kl1,
+            rtol=config.RTOL.get(str(self.mu_1.dtype)),
+            atol=config.ATOL.get(str(self.mu_1.dtype)),
+        )
+
+    def kl_divergence(self, dist1, dist2):
+        t1 = np.array(dist1._unbroadcasted_scale_tril)
+        t2 = np.array(dist2._unbroadcasted_scale_tril)
+        half_log_det_1 = np.log(t1.diagonal(axis1=-2, axis2=-1)).sum(-1)
+        half_log_det_2 = np.log(t2.diagonal(axis1=-2, axis2=-1)).sum(-1)
+        new_perm = list(range(len(t1.shape)))
+        new_perm[-1], new_perm[-2] = new_perm[-2], new_perm[-1]
+        cov_mat_1 = np.matmul(t1, t1.transpose(new_perm))
+        cov_mat_2 = np.matmul(t2, t2.transpose(new_perm))
+        expectation = (
+            np.linalg.solve(cov_mat_2, cov_mat_1)
+            .diagonal(axis1=-2, axis2=-1)
+            .sum(-1)
+        )
+        tmp = np.linalg.solve(t2, self.mu_1 - self.mu_2)
+        expectation += np.matmul(tmp.T, tmp)
+        return half_log_det_2 - half_log_det_1 + 0.5 * (expectation - 2.0)
+
+
+@parameterize.place(config.DEVICES)
+@parameterize_cls([TEST_CASE_NAME], ['MVNTestError'])
+class MVNTestError(unittest.TestCase):
+    def setUp(self):
+        paddle.disable_static(self.place)
+
+
+class TestMVNValidateArgsAndExpand(unittest.TestCase):
+    def test_mode_and_expand(self):
+        paddle.disable_static()
+        loc = paddle.to_tensor([1.0, -2.0], dtype='float32')
+        cov = paddle.to_tensor([[2.0, 0.5], [0.5, 1.5]], dtype='float32')
+        dist = MultivariateNormal(
+            loc=loc, covariance_matrix=cov, validate_args=True
+        )
+        self.assertTrue(dist._validate_args_enabled)
+        np.testing.assert_allclose(dist.mode.numpy(), loc.numpy())
+
+        expanded = dist.expand((3,))
+        self.assertTrue(expanded._validate_args_enabled)
+        self.assertEqual(expanded.batch_shape, (3,))
+        self.assertEqual(expanded.event_shape, (2,))
+        np.testing.assert_allclose(
+            expanded.mode.numpy(), np.broadcast_to(loc.numpy(), (3, 2))
+        )
+        np.testing.assert_allclose(
+            expanded.mean.numpy(), np.broadcast_to(loc.numpy(), (3, 2))
+        )
+        np.testing.assert_allclose(
+            expanded.variance.numpy(),
+            np.broadcast_to(np.diag(cov.numpy()), (3, 2)),
+        )
+
+    def test_validate_args_errors(self):
+        paddle.disable_static()
+        loc = paddle.to_tensor([0.0, 0.0], dtype='float32')
+        bad_cov = paddle.to_tensor([[1.0, 2.0], [2.0, 1.0]], dtype='float32')
+        bad_scale = paddle.to_tensor([[1.0, 0.0], [0.1, -1.0]], dtype='float32')
+        good_cov = paddle.to_tensor([[2.0, 0.5], [0.5, 1.5]], dtype='float32')
+
+        with self.assertRaises(ValueError):
+            MultivariateNormal(
+                loc=loc, covariance_matrix=bad_cov, validate_args=True
+            )
+
+        with self.assertRaises(ValueError):
+            MultivariateNormal(
+                loc=loc, scale_tril=bad_scale, validate_args=True
+            )
+
+        dist = MultivariateNormal(
+            loc=loc, covariance_matrix=good_cov, validate_args=True
+        )
+        with self.assertRaises(ValueError):
+            dist.log_prob(paddle.to_tensor([np.nan, 0.0], dtype='float32'))
+
+    def test_validate_args_additional_errors(self):
+        paddle.disable_static()
+        loc = paddle.to_tensor([0.0, 0.0], dtype='float32')
+        cov = paddle.to_tensor([[2.0, 0.5], [0.5, 1.5]], dtype='float32')
+
+        with self.assertRaises(ValueError):
+            MultivariateNormal(
+                loc=paddle.to_tensor(0.0),
+                covariance_matrix=paddle.to_tensor([[1.0]], dtype='float32'),
+            )
+
+        with self.assertRaises(ValueError):
+            MultivariateNormal(loc=loc, covariance_matrix=paddle.ones([2]))
+
+        with self.assertRaises(ValueError):
+            MultivariateNormal(loc=loc, scale_tril=paddle.ones([2]))
+
+        with self.assertRaises(ValueError):
+            MultivariateNormal(loc=loc, precision_matrix=paddle.ones([2]))
+        with self.assertRaises(ValueError):
+            MultivariateNormal(
+                loc=loc,
+                precision_matrix=paddle.to_tensor(
+                    [[1.0, 2.0], [2.0, 1.0]], dtype='float32'
+                ),
+                validate_args=True,
+            )
+
+        dist = MultivariateNormal(
+            loc=loc, covariance_matrix=cov, validate_args=True
+        )
+        with self.assertRaises(ValueError):
+            dist.log_prob(paddle.zeros([3], dtype='float32'))
+        batch_dist = MultivariateNormal(
+            loc=paddle.zeros([2, 2], dtype='float32'),
+            covariance_matrix=cov,
+            validate_args=True,
+        )
+        with self.assertRaises(ValueError):
+            batch_dist.log_prob(paddle.zeros([3, 2], dtype='float32'))
+
+    def test_validate_args_false_and_lazy_properties(self):
+        paddle.disable_static()
+        loc = paddle.to_tensor([0.0, 0.0], dtype='float32')
+        bad_scale = paddle.to_tensor([[1.0, 2.0], [0.0, 1.0]], dtype='float32')
+        dist = MultivariateNormal(
+            loc=loc, scale_tril=bad_scale, validate_args=False
+        )
+        self.assertFalse(dist._validate_args_enabled)
+
+        cov = paddle.to_tensor([[2.0, 0.5], [0.5, 1.5]], dtype='float32')
+        precision = paddle.linalg.inv(cov)
+        scale = paddle.linalg.cholesky(cov)
+
+        cov_dist = MultivariateNormal(loc=loc, covariance_matrix=cov)
+        np.testing.assert_allclose(cov_dist.scale_tril.numpy(), scale.numpy())
+        np.testing.assert_allclose(
+            cov_dist.precision_matrix.numpy(), precision.numpy(), rtol=1e-5
+        )
+
+        scale_dist = MultivariateNormal(loc=loc, scale_tril=scale)
+        scale_expanded = scale_dist.expand((3,))
+        np.testing.assert_allclose(
+            scale_expanded.scale_tril.numpy(),
+            np.broadcast_to(scale.numpy(), (3, 2, 2)),
+        )
+
+        precision_dist = MultivariateNormal(loc=loc, precision_matrix=precision)
+        precision_expanded = precision_dist.expand((3,))
+        np.testing.assert_allclose(
+            precision_dist.covariance_matrix.numpy(), cov.numpy(), rtol=1e-5
+        )
+        np.testing.assert_allclose(
+            precision_expanded.precision_matrix.numpy(),
+            np.broadcast_to(precision.numpy(), (3, 2, 2)),
+            rtol=1e-5,
+        )
+
+
+class TestMVNConstraints(unittest.TestCase):
+    def test_constraints_check(self):
+        paddle.disable_static()
+        with self.assertRaises(NotImplementedError):
+            constraint.Constraint()(paddle.ones([1], dtype='float32'))
+
+        np.testing.assert_array_equal(
+            constraint.real_vector.check(
+                paddle.to_tensor([1.0, np.nan], dtype='float32')
+            ).numpy(),
+            np.array(False),
+        )
+        np.testing.assert_array_equal(
+            constraint.real_vector.check(
+                paddle.to_tensor(1.0, dtype='float32')
+            ).numpy(),
+            np.array(False),
+        )
+
+        lower = paddle.to_tensor([[1.0, 0.0], [2.0, 3.0]], dtype='float32')
+        not_lower = paddle.to_tensor([[1.0, 2.0], [0.0, 3.0]], dtype='float32')
+        np.testing.assert_array_equal(
+            constraint.lower_triangular.check(lower).numpy(), np.array(True)
+        )
+        np.testing.assert_array_equal(
+            constraint.lower_triangular.check(not_lower).numpy(),
+            np.array(False),
+        )
+        np.testing.assert_array_equal(
+            constraint.lower_triangular.check(
+                paddle.to_tensor([1.0, 2.0], dtype='float32')
+            ).numpy(),
+            np.array(False),
+        )
+
+        bad_cholesky = paddle.to_tensor(
+            [[1.0, 0.0], [2.0, -3.0]], dtype='float32'
+        )
+        np.testing.assert_array_equal(
+            constraint.lower_cholesky.check(lower).numpy(), np.array(True)
+        )
+        np.testing.assert_array_equal(
+            constraint.lower_cholesky.check(bad_cholesky).numpy(),
+            np.array(False),
+        )
+
+        square = paddle.eye(2, dtype='float32')
+        not_square = paddle.ones([2, 3], dtype='float32')
+        not_symmetric = paddle.to_tensor(
+            [[1.0, 2.0], [0.0, 1.0]], dtype='float32'
+        )
+        not_positive_definite = paddle.to_tensor(
+            [[1.0, 2.0], [2.0, 1.0]], dtype='float32'
+        )
+        np.testing.assert_array_equal(
+            constraint.square.check(square).numpy(), np.array(True)
+        )
+        np.testing.assert_array_equal(
+            constraint.square.check(not_square).numpy(), np.array(False)
+        )
+        np.testing.assert_array_equal(
+            constraint.symmetric.check(not_symmetric).numpy(), np.array(False)
+        )
+        np.testing.assert_array_equal(
+            constraint.positive_definite.check(square).numpy(), np.array(True)
+        )
+        np.testing.assert_array_equal(
+            constraint.positive_definite.check(not_positive_definite).numpy(),
+            np.array(False),
+        )
+
+
+if __name__ == '__main__':
+    unittest.main(argv=[''], verbosity=3, exit=False)

test/distribution/test_distribution_multivariate_normal_static.py

@@ -0,0 +1,282 @@
+# Copyright (c) 2021 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
+#
+#     http://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
+import parameterize
+import scipy
+from distribution import config
+
+import paddle
+from paddle.distribution.multivariate_normal import MultivariateNormal
+
+paddle.enable_static()
+
+
+@parameterize.place(config.DEVICES)
+@parameterize.parameterize_cls(
+    (parameterize.TEST_CASE_NAME, 'loc', 'covariance_matrix'),
+    [
+        (
+            'one-batch',
+            parameterize.xrand((2,), dtype='float32', min=1, max=2),
+            np.array([[2.0, 1.0], [1.0, 2.0]]),
+        ),
+        (
+            'multi-batch',
+            parameterize.xrand((2, 3), dtype='float64', min=-2, max=-1),
+            np.array([[6.0, 2.5, 3.0], [2.5, 4.0, 5.0], [3.0, 5.0, 7.0]]),
+        ),
+    ],
+)
+class TestMVN(unittest.TestCase):
+    def setUp(self):
+        startup_program = paddle.static.Program()
+        main_program = paddle.static.Program()
+        executor = paddle.static.Executor(self.place)
+        with paddle.static.program_guard(main_program, startup_program):
+            loc = paddle.static.data('loc', self.loc.shape, self.loc.dtype)
+            covariance_matrix = paddle.static.data(
+                'covariance_matrix',
+                self.covariance_matrix.shape,
+                self.covariance_matrix.dtype,
+            )
+            dist = MultivariateNormal(
+                loc=loc, covariance_matrix=covariance_matrix
+            )
+            mean = dist.mean
+            var = dist.variance
+            entropy = dist.entropy()
+            mini_samples = dist.sample(shape=())
+            large_samples = dist.sample(shape=(10000,))
+        fetch_list = [mean, var, entropy, mini_samples, large_samples]
+        feed = {'loc': self.loc, 'covariance_matrix': self.covariance_matrix}
+
+        executor.run(startup_program)
+        [
+            self.mean,
+            self.var,
+            self.entropy,
+            self.mini_samples,
+            self.large_samples,
+        ] = executor.run(main_program, feed=feed, fetch_list=fetch_list)
+
+    def test_mean(self):
+        self.assertEqual(str(self.mean.dtype).split('.')[-1], self.loc.dtype)
+        np.testing.assert_allclose(
+            self.mean,
+            self._np_mean(),
+            rtol=config.RTOL.get(str(self.loc.dtype)),
+            atol=config.ATOL.get(str(self.loc.dtype)),
+        )
+
+    def test_variance(self):
+        self.assertEqual(str(self.var.dtype).split('.')[-1], self.loc.dtype)
+        np.testing.assert_allclose(
+            self.var,
+            self._np_variance(),
+            rtol=config.RTOL.get(str(self.loc.dtype)),
+            atol=config.ATOL.get(str(self.loc.dtype)),
+        )
+
+    def test_entropy(self):
+        self.assertEqual(str(self.entropy.dtype).split('.')[-1], self.loc.dtype)
+        np.testing.assert_allclose(
+            self.entropy,
+            self._np_entropy(),
+            rtol=config.RTOL.get(str(self.loc.dtype)),
+            atol=config.ATOL.get(str(self.loc.dtype)),
+        )
+
+    def test_sample(self):
+        self.assertEqual(
+            str(self.mini_samples.dtype).split('.')[-1], self.loc.dtype
+        )
+        sample_mean = self.large_samples.mean(axis=0)
+        sample_variance = self.large_samples.var(axis=0)
+
+        # `atol` and `rtol` refer to ``test_distribution_normal`` and ``test_distribution_lognormal``
+        np.testing.assert_allclose(sample_mean, self.mean, atol=0, rtol=0.1)
+        np.testing.assert_allclose(sample_variance, self.var, atol=0, rtol=0.1)
+
+    def _np_variance(self):
+        batch_shape = np.broadcast_shapes(
+            self.covariance_matrix.shape[:-2], self.loc.shape[:-1]
+        )
+        event_shape = self.loc.shape[-1:]
+        return np.broadcast_to(
+            np.diag(self.covariance_matrix), batch_shape + event_shape
+        )
+
+    def _np_mean(self):
+        return self.loc
+
+    def _np_entropy(self):
+        if len(self.loc.shape) <= 1:
+            return scipy.stats.multivariate_normal.entropy(
+                self.loc, self.covariance_matrix
+            )
+        else:
+            return np.apply_along_axis(
+                lambda i: scipy.stats.multivariate_normal.entropy(
+                    i, self.covariance_matrix
+                ),
+                axis=1,
+                arr=self.loc,
+            )
+
+
+@parameterize.place(config.DEVICES)
+@parameterize.parameterize_cls(
+    (parameterize.TEST_CASE_NAME, 'loc', 'covariance_matrix', 'value'),
+    [
+        (
+            'value-same-shape',
+            parameterize.xrand((2,), dtype='float32', min=-2, max=2),
+            np.array([[2.0, 1.0], [1.0, 2.0]]),
+            parameterize.xrand((2,), dtype='float32', min=-5, max=5),
+        ),
+        (
+            'value-broadcast-shape',
+            parameterize.xrand((2,), dtype='float64', min=-2, max=2),
+            np.array([[2.0, 1.0], [1.0, 2.0]]),
+            parameterize.xrand((3, 2), dtype='float64', min=-5, max=5),
+        ),
+    ],
+)
+class TestMVNProbs(unittest.TestCase):
+    def setUp(self):
+        startup_program = paddle.static.Program()
+        main_program = paddle.static.Program()
+        executor = paddle.static.Executor(self.place)
+
+        with paddle.static.program_guard(main_program, startup_program):
+            loc = paddle.static.data('loc', self.loc.shape, self.loc.dtype)
+            covariance_matrix = paddle.static.data(
+                'covariance_matrix',
+                self.covariance_matrix.shape,
+                self.covariance_matrix.dtype,
+            )
+            value = paddle.static.data(
+                'value', self.value.shape, self.value.dtype
+            )
+            dist = MultivariateNormal(
+                loc=loc, covariance_matrix=covariance_matrix
+            )
+            pmf = dist.prob(value)
+        feed = {
+            'loc': self.loc,
+            'covariance_matrix': self.covariance_matrix,
+            'value': self.value,
+        }
+        fetch_list = [pmf]
+
+        executor.run(startup_program)
+        [self.pmf] = executor.run(
+            main_program, feed=feed, fetch_list=fetch_list
+        )
+
+    def test_prob(self):
+        if len(self.value.shape) <= 1:
+            scipy_pdf = scipy.stats.multivariate_normal.pdf(
+                self.value, self.loc, self.covariance_matrix
+            )
+        else:
+            scipy_pdf = np.apply_along_axis(
+                lambda i: scipy.stats.multivariate_normal.pdf(
+                    i, self.loc, self.covariance_matrix
+                ),
+                axis=1,
+                arr=self.value,
+            )
+        np.testing.assert_allclose(
+            self.pmf,
+            scipy_pdf,
+            rtol=config.RTOL.get(str(self.loc.dtype)),
+            atol=config.ATOL.get(str(self.loc.dtype)),
+        )
+
+
+@parameterize.place(config.DEVICES)
+@parameterize.parameterize_cls(
+    (parameterize.TEST_CASE_NAME, 'mu_1', 'sig_1', 'mu_2', 'sig_2'),
+    [
+        (
+            'one-batch',
+            parameterize.xrand((2,), dtype='float32', min=-2, max=2),
+            np.array([[2.0, 1.0], [1.0, 2.0]]).astype('float32'),
+            parameterize.xrand((2,), dtype='float32', min=-2, max=2),
+            np.array([[3.0, 2.0], [2.0, 3.0]]).astype('float32'),
+        )
+    ],
+)
+class TestMVNKL(unittest.TestCase):
+    def setUp(self):
+        startup_program = paddle.static.Program()
+        main_program = paddle.static.Program()
+        executor = paddle.static.Executor(self.place)
+
+        with paddle.static.program_guard(main_program, startup_program):
+            mu_1 = paddle.static.data('mu_1', self.mu_1.shape)
+            sig_1 = paddle.static.data('sig_1', self.sig_1.shape)
+            mu_2 = paddle.static.data('mu_2', self.mu_2.shape)
+            sig_2 = paddle.static.data('sig_2', self.sig_2.shape)
+            dist1 = MultivariateNormal(loc=mu_1, covariance_matrix=sig_1)
+            dist2 = MultivariateNormal(loc=mu_2, covariance_matrix=sig_2)
+            kl_dist1_dist2 = dist1.kl_divergence(dist2)
+        feed = {
+            'mu_1': self.mu_1,
+            'sig_1': self.sig_1,
+            'mu_2': self.mu_2,
+            'sig_2': self.sig_2,
+        }
+        fetch_list = [kl_dist1_dist2]
+
+        executor.run(startup_program)
+        [self.kl_dist1_dist2] = executor.run(
+            main_program, feed=feed, fetch_list=fetch_list
+        )
+
+    def test_kl_divergence(self):
+        kl0 = self.kl_dist1_dist2
+        kl1 = self.kl_divergence()
+        batch_shape = np.broadcast_shapes(
+            self.sig_1.shape[:-2], self.mu_1.shape[:-1]
+        )
+        self.assertEqual(tuple(kl0.shape), batch_shape)
+        self.assertEqual(tuple(kl1.shape), batch_shape)
+        np.testing.assert_allclose(kl0, kl1, rtol=0.1, atol=0.1)
+
+    def kl_divergence(self):
+        t1 = np.array(np.linalg.cholesky(self.sig_1))
+        t2 = np.array(np.linalg.cholesky(self.sig_2))
+        half_log_det_1 = np.log(t1.diagonal(axis1=-2, axis2=-1)).sum(-1)
+        half_log_det_2 = np.log(t2.diagonal(axis1=-2, axis2=-1)).sum(-1)
+        new_perm = list(range(len(t1.shape)))
+        new_perm[-1], new_perm[-2] = new_perm[-2], new_perm[-1]
+        cov_mat_1 = np.matmul(t1, t1.transpose(new_perm))
+        cov_mat_2 = np.matmul(t2, t2.transpose(new_perm))
+        expectation = (
+            np.linalg.solve(cov_mat_2, cov_mat_1)
+            .diagonal(axis1=-2, axis2=-1)
+            .sum(-1)
+        )
+        tmp = np.linalg.solve(t2, self.mu_1 - self.mu_2)
+        expectation += np.matmul(tmp.T, tmp)
+        return half_log_det_2 - half_log_det_1 + 0.5 * (expectation - 2.0)
+
+
+if __name__ == '__main__':
+    unittest.main(argv=[''], verbosity=3, exit=False)

test/distribution/test_distribution_poisson.py

@@ -0,0 +1,205 @@
+# Copyright (c) 2021 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
+#
+#     http://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
+import parameterize
+import scipy.stats
+from distribution import config
+
+import paddle
+from paddle.distribution import Poisson
+
+
+@parameterize.place(config.DEVICES)
+@parameterize.parameterize_cls(
+    (parameterize.TEST_CASE_NAME, 'rate'),
+    [
+        ('zero-dim', np.array(100.0).astype('float64')),
+        ('one-dim', np.array([100.0]).astype('float64')),
+        # boundary case and extreme case (`scipy.stats.poisson.entropy` cannot converge for very extreme cases such as rate=10000.0)
+        ('multi-dim', np.array([0.0, 1000.0]).astype('float32')),
+    ],
+)
+class TestPoisson(unittest.TestCase):
+    def setUp(self):
+        self._dist = Poisson(rate=paddle.to_tensor(self.rate))
+
+    def test_mean(self):
+        mean = self._dist.mean
+        self.assertEqual(mean.numpy().dtype, self.rate.dtype)
+        np.testing.assert_allclose(
+            mean,
+            scipy.stats.poisson.mean(self.rate),
+            rtol=config.RTOL.get(str(self.rate.dtype)),
+            atol=config.ATOL.get(str(self.rate.dtype)),
+        )
+
+    def test_variance(self):
+        var = self._dist.variance
+        self.assertEqual(var.numpy().dtype, self.rate.dtype)
+        np.testing.assert_allclose(
+            var,
+            scipy.stats.poisson.var(self.rate),
+            rtol=config.RTOL.get(str(self.rate.dtype)),
+            atol=config.ATOL.get(str(self.rate.dtype)),
+        )
+
+    def test_entropy(self):
+        entropy = self._dist.entropy()
+        self.assertEqual(entropy.numpy().dtype, self.rate.dtype)
+        np.testing.assert_allclose(
+            entropy,
+            scipy.stats.poisson.entropy(self.rate),
+            rtol=config.RTOL.get(str(self.rate.dtype)),
+            atol=config.ATOL.get(str(self.rate.dtype)),
+        )
+
+    def test_sample(self):
+        sample_shape = ()
+        samples = self._dist.sample(sample_shape)
+        self.assertEqual(samples.numpy().dtype, self.rate.dtype)
+        self.assertEqual(
+            tuple(samples.shape),
+            sample_shape + self._dist.batch_shape + self._dist.event_shape,
+        )
+
+        sample_shape = (5000,)
+        samples = self._dist.sample(sample_shape)
+        sample_mean = samples.mean(axis=0)
+        sample_variance = samples.var(axis=0)
+
+        np.testing.assert_allclose(
+            sample_mean, self._dist.mean, atol=0, rtol=0.20
+        )
+        np.testing.assert_allclose(
+            sample_variance, self._dist.variance, atol=0, rtol=0.20
+        )
+
+
+@parameterize.place(config.DEVICES)
+@parameterize.parameterize_cls(
+    (parameterize.TEST_CASE_NAME, 'rate', 'value'),
+    [
+        (
+            'value-same-shape',
+            np.array(1000).astype('float32'),
+            np.array(1100).astype('float32'),
+        ),
+        (
+            'value-broadcast-shape',
+            np.array(10).astype('float64'),
+            np.array([2.0, 3.0, 5.0, 10.0, 20.0]).astype('float64'),
+        ),
+    ],
+)
+class TestPoissonProbs(unittest.TestCase):
+    def setUp(self):
+        self._dist = Poisson(rate=paddle.to_tensor(self.rate))
+
+    def test_prob(self):
+        np.testing.assert_allclose(
+            self._dist.prob(paddle.to_tensor(self.value)),
+            scipy.stats.poisson.pmf(self.value, self.rate),
+            rtol=config.RTOL.get(str(self.rate.dtype)),
+            atol=config.ATOL.get(str(self.rate.dtype)),
+        )
+
+    def test_log_prob(self):
+        np.testing.assert_allclose(
+            self._dist.log_prob(paddle.to_tensor(self.value)),
+            scipy.stats.poisson.logpmf(self.value, self.rate),
+            rtol=config.RTOL.get(str(self.rate.dtype)),
+            atol=config.ATOL.get(str(self.rate.dtype)),
+        )
+
+
+@parameterize.place(config.DEVICES)
+@parameterize.parameterize_cls(
+    (parameterize.TEST_CASE_NAME, 'rate_1', 'rate_2'),
+    [
+        (
+            'zero-dim',
+            parameterize.xrand((1,), min=1, max=20)
+            .astype('int32')
+            .astype('float64')
+            .reshape([]),
+            parameterize.xrand((1,), min=1, max=20)
+            .astype('int32')
+            .astype('float64')
+            .reshape([]),
+        ),
+        (
+            'one-dim',
+            parameterize.xrand((1,), min=1, max=20)
+            .astype('int32')
+            .astype('float64'),
+            parameterize.xrand((1,), min=1, max=20)
+            .astype('int32')
+            .astype('float64'),
+        ),
+        (
+            'multi-dim',
+            parameterize.xrand((5, 3), min=1, max=20)
+            .astype('int32')
+            .astype('float32'),
+            parameterize.xrand((5, 3), min=1, max=20)
+            .astype('int32')
+            .astype('float32'),
+        ),
+    ],
+)
+class TestPoissonKL(unittest.TestCase):
+    def setUp(self):
+        self._dist1 = Poisson(rate=paddle.to_tensor(self.rate_1))
+        self._dist2 = Poisson(rate=paddle.to_tensor(self.rate_2))
+
+    def test_kl_divergence(self):
+        kl0 = self._dist1.kl_divergence(self._dist2)
+        kl1 = self.kl_divergence_scipy()
+
+        self.assertEqual(tuple(kl0.shape), self._dist1.batch_shape)
+        self.assertEqual(tuple(kl1.shape), self._dist1.batch_shape)
+        np.testing.assert_allclose(
+            kl0,
+            kl1,
+            rtol=config.RTOL.get(str(self.rate_1.dtype)),
+            atol=config.ATOL.get(str(self.rate_1.dtype)),
+        )
+
+    def kl_divergence_scipy(self):
+        rate_max = np.max(np.maximum(self.rate_1, self.rate_2))
+        rate_min = np.min(np.minimum(self.rate_1, self.rate_2))
+        support_max = self.enumerate_bounded_support(rate_max)
+        support_min = self.enumerate_bounded_support(rate_min)
+        a_min = np.min(support_min)
+        a_max = np.max(support_max)
+        common_support = np.arange(
+            a_min, a_max, dtype=self.rate_1.dtype
+        ).reshape((-1,) + (1,) * len(self.rate_1.shape))
+        log_prob_1 = scipy.stats.poisson.logpmf(common_support, self.rate_1)
+        log_prob_2 = scipy.stats.poisson.logpmf(common_support, self.rate_2)
+        return (np.exp(log_prob_1) * (log_prob_1 - log_prob_2)).sum(0)
+
+    def enumerate_bounded_support(self, rate):
+        s = np.sqrt(rate)
+        upper = int(rate + 30 * s)
+        lower = int(np.clip(rate - 30 * s, a_min=0, a_max=rate))
+        values = np.arange(lower, upper, dtype=self.rate_1.dtype)
+        return values
+
+
+if __name__ == '__main__':
+    unittest.main()

test/distribution/test_distribution_poisson_static.py

@@ -0,0 +1,246 @@
+# Copyright (c) 2021 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
+#
+#     http://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
+import parameterize
+import scipy.stats
+from distribution import config
+
+import paddle
+from paddle.distribution import Poisson
+
+paddle.enable_static()
+
+
+paddle.enable_static()
+
+
+@parameterize.place(config.DEVICES)
+@parameterize.parameterize_cls(
+    (parameterize.TEST_CASE_NAME, 'rate'),
+    [
+        ('zero-dim', np.array(1000.0).astype('float32')),
+        ('one-dim', np.array([1000.0]).astype('float32')),
+        (
+            'multi-dim',
+            parameterize.xrand((2,), min=1, max=20)
+            .astype('int32')
+            .astype('float64'),
+        ),
+    ],
+)
+class TestPoisson(unittest.TestCase):
+    def setUp(self):
+        startup_program = paddle.static.Program()
+        main_program = paddle.static.Program()
+        executor = paddle.static.Executor(self.place)
+        with paddle.static.program_guard(main_program, startup_program):
+            rate = paddle.static.data('rate', self.rate.shape, self.rate.dtype)
+            dist = Poisson(rate)
+            mean = dist.mean
+            var = dist.variance
+            entropy = dist.entropy()
+            mini_samples = dist.sample(shape=())
+            large_samples = dist.sample(shape=(1000,))
+        fetch_list = [mean, var, entropy, mini_samples, large_samples]
+        feed = {'rate': self.rate}
+
+        executor.run(startup_program)
+        [
+            self.mean,
+            self.var,
+            self.entropy,
+            self.mini_samples,
+            self.large_samples,
+        ] = executor.run(main_program, feed=feed, fetch_list=fetch_list)
+
+    def test_mean(self):
+        self.assertEqual(str(self.mean.dtype).split('.')[-1], self.rate.dtype)
+        np.testing.assert_allclose(
+            self.mean,
+            self._np_mean(),
+            rtol=config.RTOL.get(str(self.rate.dtype)),
+            atol=config.ATOL.get(str(self.rate.dtype)),
+        )
+
+    def test_variance(self):
+        self.assertEqual(str(self.var.dtype).split('.')[-1], self.rate.dtype)
+        np.testing.assert_allclose(
+            self.var,
+            self._np_variance(),
+            rtol=config.RTOL.get(str(self.rate.dtype)),
+            atol=config.ATOL.get(str(self.rate.dtype)),
+        )
+
+    def test_entropy(self):
+        self.assertEqual(
+            str(self.entropy.dtype).split('.')[-1], self.rate.dtype
+        )
+        np.testing.assert_allclose(
+            self.entropy,
+            self._np_entropy(),
+            rtol=config.RTOL.get(str(self.rate.dtype)),
+            atol=config.ATOL.get(str(self.rate.dtype)),
+        )
+
+    def test_sample(self):
+        self.assertEqual(
+            str(self.mini_samples.dtype).split('.')[-1], self.rate.dtype
+        )
+        sample_mean = self.large_samples.mean(axis=0)
+        sample_variance = self.large_samples.var(axis=0)
+        np.testing.assert_allclose(sample_mean, self.mean, atol=0, rtol=0.20)
+        np.testing.assert_allclose(sample_variance, self.var, atol=0, rtol=0.20)
+
+    def _np_variance(self):
+        return self.rate
+
+    def _np_mean(self):
+        return self.rate
+
+    def _np_entropy(self):
+        return scipy.stats.poisson.entropy(self.rate)
+
+
+@parameterize.place(config.DEVICES)
+@parameterize.parameterize_cls(
+    (parameterize.TEST_CASE_NAME, 'rate', 'value'),
+    [
+        (
+            'value-same-shape',
+            np.array(1000).astype('float32'),
+            np.array(1100).astype('float32'),
+        ),
+        (
+            'value-broadcast-shape',
+            np.array(10).astype('float64'),
+            np.array([2.0, 3.0]).astype('float64'),
+        ),
+    ],
+)
+class TestPoissonProbs(unittest.TestCase):
+    def setUp(self):
+        startup_program = paddle.static.Program()
+        main_program = paddle.static.Program()
+        executor = paddle.static.Executor(self.place)
+
+        with paddle.static.program_guard(main_program, startup_program):
+            rate = paddle.static.data('rate', self.rate.shape, self.rate.dtype)
+            value = paddle.static.data(
+                'value', self.value.shape, self.value.dtype
+            )
+            dist = Poisson(rate)
+            pmf = dist.prob(value)
+        feed = {'rate': self.rate, 'value': self.value}
+        fetch_list = [pmf]
+
+        executor.run(startup_program)
+        [self.pmf] = executor.run(
+            main_program, feed=feed, fetch_list=fetch_list
+        )
+
+    def test_prob(self):
+        np.testing.assert_allclose(
+            self.pmf,
+            scipy.stats.poisson.pmf(self.value, self.rate),
+            rtol=config.RTOL.get(str(self.rate.dtype)),
+            atol=config.ATOL.get(str(self.rate.dtype)),
+        )
+
+
+@parameterize.place(config.DEVICES)
+@parameterize.parameterize_cls(
+    (parameterize.TEST_CASE_NAME, 'rate_1', 'rate_2'),
+    [
+        (
+            'zero-dim',
+            parameterize.xrand((1,), min=1, max=20)
+            .astype('int32')
+            .astype('float32')
+            .reshape([]),
+            parameterize.xrand((1,), min=1, max=20)
+            .astype('int32')
+            .astype('float32')
+            .reshape([]),
+        ),
+        (
+            'multi-dim',
+            parameterize.xrand((2, 3), min=1, max=20)
+            .astype('int32')
+            .astype('float32'),
+            parameterize.xrand((2, 3), min=1, max=20)
+            .astype('int32')
+            .astype('float32'),
+        ),
+    ],
+)
+class TestPoissonKL(unittest.TestCase):
+    def setUp(self):
+        startup_program = paddle.static.Program()
+        main_program = paddle.static.Program()
+        executor = paddle.static.Executor(self.place)
+
+        with paddle.static.program_guard(main_program, startup_program):
+            rate_1 = paddle.static.data('rate_1', self.rate_1.shape)
+            rate_2 = paddle.static.data('rate_2', self.rate_2.shape)
+            dist1 = Poisson(rate_1)
+            dist2 = Poisson(rate_2)
+            kl_dist1_dist2 = dist1.kl_divergence(dist2)
+        feed = {'rate_1': self.rate_1, 'rate_2': self.rate_2}
+        fetch_list = [kl_dist1_dist2]
+
+        executor.run(startup_program)
+        [self.kl_dist1_dist2] = executor.run(
+            main_program, feed=feed, fetch_list=fetch_list
+        )
+
+    def test_kl_divergence(self):
+        kl0 = self.kl_dist1_dist2
+        kl1 = self.kl_divergence_scipy()
+
+        self.assertEqual(tuple(kl0.shape), self.rate_1.shape)
+        self.assertEqual(tuple(kl1.shape), self.rate_1.shape)
+        np.testing.assert_allclose(
+            kl0,
+            kl1,
+            rtol=config.RTOL.get(str(self.rate_1.dtype)),
+            atol=config.ATOL.get(str(self.rate_1.dtype)),
+        )
+
+    def kl_divergence_scipy(self):
+        rate_max = np.max(np.maximum(self.rate_1, self.rate_2))
+        rate_min = np.min(np.minimum(self.rate_1, self.rate_2))
+        support_max = self.enumerate_bounded_support(rate_max)
+        support_min = self.enumerate_bounded_support(rate_min)
+        a_min = np.min(support_min)
+        a_max = np.max(support_max)
+        common_support = np.arange(
+            a_min, a_max, dtype=self.rate_1.dtype
+        ).reshape((-1,) + (1,) * len(self.rate_1.shape))
+        log_prob_1 = scipy.stats.poisson.logpmf(common_support, self.rate_1)
+        log_prob_2 = scipy.stats.poisson.logpmf(common_support, self.rate_2)
+        return (np.exp(log_prob_1) * (log_prob_1 - log_prob_2)).sum(0)
+
+    def enumerate_bounded_support(self, rate):
+        s = np.sqrt(rate)
+        upper = int(rate + 30 * s)
+        lower = int(np.clip(rate - 30 * s, a_min=0, a_max=rate))
+        values = np.arange(lower, upper, dtype=self.rate_1.dtype)
+        return values
+
+
+if __name__ == '__main__':
+    unittest.main()

test/distribution/test_distribution_student_t.py

@@ -0,0 +1,340 @@
+# Copyright (c) 2024 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
+#
+#     http://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
+import parameterize
+import scipy.stats
+from distribution import config
+from parameterize import (
+    TEST_CASE_NAME,
+    parameterize_cls,
+    parameterize_func,
+)
+
+import paddle
+from paddle.distribution.student_t import StudentT
+
+
+@parameterize.place(config.DEVICES)
+@parameterize.parameterize_cls(
+    (parameterize.TEST_CASE_NAME, 'df', 'loc', 'scale'),
+    [
+        (
+            'one-dim',
+            10.0,
+            1.0,
+            2.0,
+        ),
+        (
+            'multi-dim',
+            parameterize.xrand((2, 1), dtype='float32', min=4, max=30),
+            parameterize.xrand((2, 3), dtype='float32', min=1, max=10),
+            parameterize.xrand((2, 3), dtype='float32', min=0.1, max=3),
+        ),
+        (
+            'multi-dim2',
+            parameterize.xrand((2, 1), dtype='float64', min=4, max=30),
+            parameterize.xrand((2, 3), dtype='float64', min=-10, max=-1),
+            parameterize.xrand((2, 3), dtype='float64', min=0.1, max=3),
+        ),
+    ],
+)
+class TestStudentT(unittest.TestCase):
+    def setUp(self):
+        df = (
+            self.df if isinstance(self.df, float) else paddle.to_tensor(self.df)
+        )
+        loc = (
+            self.loc
+            if isinstance(self.loc, float)
+            else paddle.to_tensor(self.loc)
+        )
+        scale = (
+            self.scale
+            if isinstance(self.scale, float)
+            else paddle.to_tensor(self.scale)
+        )
+        self._dist = StudentT(df, loc, scale)
+
+    def test_mean(self):
+        mean = self._dist.mean
+        target_dtype = (
+            "float32" if isinstance(self.df, float) else self.df.dtype
+        )
+        self.assertEqual(mean.numpy().dtype, target_dtype)
+        np.testing.assert_allclose(
+            mean,
+            self._np_mean(),
+            rtol=config.RTOL.get(str(target_dtype)),
+            atol=config.ATOL.get(str(target_dtype)),
+        )
+
+    def test_variance(self):
+        var = self._dist.variance
+        target_dtype = (
+            "float32" if isinstance(self.df, float) else self.df.dtype
+        )
+        self.assertEqual(var.numpy().dtype, target_dtype)
+        np.testing.assert_allclose(
+            var,
+            self._np_variance(),
+            rtol=config.RTOL.get(str(target_dtype)),
+            atol=config.ATOL.get(str(target_dtype)),
+        )
+
+    def test_entropy(self):
+        entropy = self._dist.entropy()
+        target_dtype = (
+            "float32" if isinstance(self.df, float) else self.df.dtype
+        )
+        self.assertEqual(entropy.numpy().dtype, target_dtype)
+        np.testing.assert_allclose(
+            entropy,
+            self._np_entropy(),
+            rtol=config.RTOL.get(str(target_dtype)),
+            atol=config.ATOL.get(str(target_dtype)),
+        )
+
+    def test_sample(self):
+        sample_shape = ()
+        samples = self._dist.sample(sample_shape)
+        self.assertEqual(
+            tuple(samples.shape),
+            sample_shape + self._dist.batch_shape + self._dist.event_shape,
+        )
+
+        sample_shape = (10000,)
+        samples = self._dist.sample(sample_shape)
+        sample_mean = samples.mean(axis=0)
+        sample_variance = samples.var(axis=0)
+
+        # Tolerance value 0.1 is empirical value which is consistent with
+        # TensorFlow
+        np.testing.assert_allclose(
+            sample_mean, self._dist.mean, atol=0, rtol=0.10
+        )
+        # Tolerance value 0.1 is empirical value which is consistent with
+        # TensorFlow
+        np.testing.assert_allclose(
+            sample_variance, self._dist.variance, atol=0, rtol=0.10
+        )
+
+    def _np_variance(self):
+        if isinstance(self.df, np.ndarray) and self.df.dtype == np.float32:
+            df = self.df.astype("float64")
+        else:
+            df = self.df
+        if isinstance(self.loc, np.ndarray) and self.loc.dtype == np.float32:
+            loc = self.loc.astype("float64")
+        else:
+            loc = self.loc
+        if (
+            isinstance(self.scale, np.ndarray)
+            and self.scale.dtype == np.float32
+        ):
+            scale = self.scale.astype("float64")
+        else:
+            scale = self.scale
+        return scipy.stats.t.var(df, loc, scale)
+
+    def _np_mean(self):
+        if isinstance(self.df, np.ndarray) and self.df.dtype == np.float32:
+            df = self.df.astype("float64")
+        else:
+            df = self.df
+        if isinstance(self.loc, np.ndarray) and self.loc.dtype == np.float32:
+            loc = self.loc.astype("float64")
+        else:
+            loc = self.loc
+        if (
+            isinstance(self.scale, np.ndarray)
+            and self.scale.dtype == np.float32
+        ):
+            scale = self.scale.astype("float64")
+        else:
+            scale = self.scale
+        return scipy.stats.t.mean(df, loc, scale)
+
+    def _np_entropy(self):
+        if isinstance(self.df, np.ndarray) and self.df.dtype == np.float32:
+            df = self.df.astype("float64")
+        else:
+            df = self.df
+        if isinstance(self.loc, np.ndarray) and self.loc.dtype == np.float32:
+            loc = self.loc.astype("float64")
+        else:
+            loc = self.loc
+        if (
+            isinstance(self.scale, np.ndarray)
+            and self.scale.dtype == np.float32
+        ):
+            scale = self.scale.astype("float64")
+        else:
+            scale = self.scale
+        return scipy.stats.t.entropy(df, loc, scale)
+
+
+@parameterize.place(config.DEVICES)
+@parameterize.parameterize_cls(
+    (parameterize.TEST_CASE_NAME, 'df', 'loc', 'scale'),
+    [
+        (
+            'float-tensor',
+            10.0,
+            paddle.to_tensor(1.0),
+            2.0,
+        ),
+        (
+            'float-tensor1',
+            10.0,
+            parameterize.xrand((2, 3), dtype='float32', min=1, max=10),
+            2.0,
+        ),
+        (
+            'float-tensor2',
+            parameterize.xrand((2, 1), dtype='float64', min=4, max=30),
+            parameterize.xrand((2, 3), dtype='float64', min=1, max=10),
+            2.0,
+        ),
+        (
+            'float-tensor3',
+            parameterize.xrand((2, 1), dtype='float64', min=4, max=30),
+            1.0,
+            parameterize.xrand((2, 1), dtype='float64', min=0.1, max=3),
+        ),
+        (
+            'float-tensor4',
+            5.0,
+            parameterize.xrand((2, 1), dtype='float32', min=-1, max=-10),
+            parameterize.xrand((2, 3), dtype='float32', min=0.1, max=3),
+        ),
+    ],
+)
+class TestStudentT2(TestStudentT):
+    def setUp(self):
+        self._dist = StudentT(self.df, self.loc, self.scale)
+
+
+@parameterize.place(config.DEVICES)
+@parameterize.parameterize_cls(
+    (parameterize.TEST_CASE_NAME, 'df', 'loc', 'scale', 'value'),
+    [
+        (
+            'one-dim',
+            10.0,
+            0.0,
+            1.0,
+            np.array(3.3).astype("float32"),
+        ),
+        (
+            'value-broadcast-shape',
+            parameterize.xrand((2, 1), dtype='float64', min=4, max=30),
+            parameterize.xrand((2, 1), dtype='float64', min=-10, max=10),
+            parameterize.xrand((2, 1), dtype='float64', min=0.1, max=5),
+            parameterize.xrand((2, 4), dtype='float64', min=-10, max=10),
+        ),
+    ],
+)
+class TestStudentTProbs(unittest.TestCase):
+    def setUp(self):
+        df = (
+            self.df if isinstance(self.df, float) else paddle.to_tensor(self.df)
+        )
+        loc = (
+            self.loc
+            if isinstance(self.loc, float)
+            else paddle.to_tensor(self.loc)
+        )
+        scale = (
+            self.scale
+            if isinstance(self.scale, float)
+            else paddle.to_tensor(self.scale)
+        )
+        self._dist = StudentT(df, loc, scale)
+
+    def test_prob(self):
+        target_dtype = (
+            "float32" if isinstance(self.df, float) else self.df.dtype
+        )
+        np.testing.assert_allclose(
+            self._dist.prob(paddle.to_tensor(self.value)),
+            scipy.stats.t.pdf(self.value, self.df, self.loc, self.scale),
+            rtol=config.RTOL.get(str(target_dtype)),
+            atol=config.ATOL.get(str(target_dtype)),
+        )
+
+    def test_log_prob(self):
+        target_dtype = (
+            "float32" if isinstance(self.df, float) else self.df.dtype
+        )
+        np.testing.assert_allclose(
+            self._dist.log_prob(paddle.to_tensor(self.value)),
+            scipy.stats.t.logpdf(self.value, self.df, self.loc, self.scale),
+            rtol=config.RTOL.get(str(target_dtype)),
+            atol=config.ATOL.get(str(target_dtype)),
+        )
+
+
+@parameterize.place(config.DEVICES)
+@parameterize.parameterize_cls(
+    (parameterize.TEST_CASE_NAME, 'df', 'loc', 'scale', 'value'),
+    [
+        (
+            'float-tensor1',
+            10.0,
+            parameterize.xrand((2, 1), dtype='float32', min=-10, max=10),
+            1.0,
+            np.array(3.3).astype("float32"),
+        ),
+        (
+            'float-tensor2',
+            parameterize.xrand((2, 1), dtype='float64', min=4, max=30),
+            1.0,
+            parameterize.xrand((2, 1), dtype='float64', min=0.1, max=5),
+            parameterize.xrand((2, 4), dtype='float64', min=-10, max=10),
+        ),
+    ],
+)
+class TestStudentTProbs2(TestStudentTProbs):
+    def setUp(self):
+        self._dist = StudentT(self.df, self.loc, self.scale)
+
+
+@parameterize.place(config.DEVICES)
+@parameterize_cls([TEST_CASE_NAME], ['StudentTTestError'])
+class StudentTTestError(unittest.TestCase):
+    def setUp(self):
+        paddle.disable_static(self.place)
+
+    @parameterize_func(
+        [
+            (-5.0, 0.0, 1.0, ValueError),  # negative df
+            (5.0, 0.0, -1.0, ValueError),  # negative scale
+        ]
+    )
+    def test_bad_parameter(self, df, loc, scale, error):
+        with paddle.base.dygraph.guard(self.place):
+            self.assertRaises(error, StudentT, df, loc, scale)
+
+    @parameterize_func([(10,)])  # not sequence object sample shape
+    def test_bad_sample_shape(self, shape):
+        with paddle.base.dygraph.guard(self.place):
+            t = StudentT(5.0, 0.0, 1.0)
+            self.assertRaises(TypeError, t.sample, shape)
+
+
+if __name__ == '__main__':
+    unittest.main()

test/distribution/test_distribution_transformed_distribution.py

@@ -0,0 +1,79 @@
+# Copyright (c) 2021 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
+#
+#     http://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
+import parameterize as param
+from distribution import config
+
+import paddle
+
+
+@param.place(config.DEVICES)
+@param.param_cls(
+    (param.TEST_CASE_NAME, 'base', 'transforms'),
+    [
+        (
+            'base_normal',
+            paddle.distribution.Normal(0.0, 1.0),
+            [paddle.distribution.ExpTransform()],
+        )
+    ],
+)
+class TestIndependent(unittest.TestCase):
+    def setUp(self):
+        self._t = paddle.distribution.TransformedDistribution(
+            self.base, self.transforms
+        )
+
+    def _np_sum_rightmost(self, value, n):
+        return np.sum(value, tuple(range(-n, 0))) if n > 0 else value
+
+    def test_log_prob(self):
+        value = paddle.to_tensor([0.5])
+        np.testing.assert_allclose(
+            self.simple_log_prob(value, self.base, self.transforms),
+            self._t.log_prob(value),
+            rtol=config.RTOL.get(str(value.numpy().dtype)),
+            atol=config.ATOL.get(str(value.numpy().dtype)),
+        )
+
+    def simple_log_prob(self, value, base, transforms):
+        log_prob = 0.0
+        y = value
+        for t in reversed(transforms):
+            x = t.inverse(y)
+            log_prob = log_prob - t.forward_log_det_jacobian(x)
+            y = x
+        log_prob += base.log_prob(y)
+        return log_prob
+
+    # TODO(cxxly): Add Kolmogorov-Smirnov test for sample result.
+    def test_sample(self):
+        shape = [5, 10, 8]
+        expected_shape = (5, 10, 8)
+        data = self._t.sample(shape)
+        self.assertEqual(tuple(data.shape), expected_shape)
+        self.assertEqual(data.dtype, self.base.loc.dtype)
+
+    def test_rsample(self):
+        shape = [5, 10, 8]
+        expected_shape = (5, 10, 8)
+        data = self._t.rsample(shape)
+        self.assertEqual(tuple(data.shape), expected_shape)
+        self.assertEqual(data.dtype, self.base.loc.dtype)
+
+
+if __name__ == '__main__':
+    unittest.main()

test/distribution/test_distribution_transformed_distribution_static.py

@@ -0,0 +1,97 @@
+# Copyright (c) 2021 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
+#
+#     http://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
+import parameterize as param
+from distribution import config
+
+import paddle
+
+paddle.enable_static()
+
+
+@param.place(config.DEVICES)
+@param.param_cls(
+    (param.TEST_CASE_NAME, 'base', 'transforms'),
+    [
+        (
+            'base_normal',
+            paddle.distribution.Normal,
+            [paddle.distribution.ExpTransform()],
+        )
+    ],
+)
+class TestIndependent(unittest.TestCase):
+    def setUp(self):
+        value = np.array([0.5])
+        loc = np.array([0.0])
+        scale = np.array([1.0])
+        shape = [5, 10, 8]
+        self.dtype = value.dtype
+        exe = paddle.static.Executor()
+        sp = paddle.static.Program()
+        mp = paddle.static.Program()
+        with paddle.static.program_guard(mp, sp):
+            static_value = paddle.static.data('value', value.shape, value.dtype)
+            static_loc = paddle.static.data('loc', loc.shape, loc.dtype)
+            static_scale = paddle.static.data('scale', scale.shape, scale.dtype)
+            self.base = self.base(static_loc, static_scale)
+            self._t = paddle.distribution.TransformedDistribution(
+                self.base, self.transforms
+            )
+            actual_log_prob = self._t.log_prob(static_value)
+            expected_log_prob = self.transformed_log_prob(
+                static_value, self.base, self.transforms
+            )
+            sample_data = self._t.sample(shape)
+
+        exe.run(sp)
+        [
+            self.actual_log_prob,
+            self.expected_log_prob,
+            self.sample_data,
+        ] = exe.run(
+            mp,
+            feed={'value': value, 'loc': loc, 'scale': scale},
+            fetch_list=[actual_log_prob, expected_log_prob, sample_data],
+        )
+
+    def test_log_prob(self):
+        np.testing.assert_allclose(
+            self.actual_log_prob,
+            self.expected_log_prob,
+            rtol=config.RTOL.get(str(self.dtype)),
+            atol=config.ATOL.get(str(self.dtype)),
+        )
+
+    def transformed_log_prob(self, value, base, transforms):
+        log_prob = 0.0
+        y = value
+        for t in reversed(transforms):
+            x = t.inverse(y)
+            log_prob = log_prob - t.forward_log_det_jacobian(x)
+            y = x
+        log_prob += base.log_prob(y)
+        return log_prob
+
+    # TODO(cxxly): Add Kolmogorov-Smirnov test for sample result.
+    def test_sample(self):
+        expected_shape = (5, 10, 8, 1)
+        self.assertEqual(tuple(self.sample_data.shape), expected_shape)
+        self.assertEqual(self.sample_data.dtype, self.dtype)
+
+
+if __name__ == '__main__':
+    unittest.main()

test/distribution/test_distribution_uniform.py

@@ -0,0 +1,386 @@
+# Copyright (c) 2021 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
+#
+#     http://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 test_distribution import DistributionNumpy
+
+import paddle
+from paddle import base
+from paddle.distribution import Uniform
+
+np.random.seed(2022)
+
+
+class UniformNumpy(DistributionNumpy):
+    def __init__(self, low, high):
+        self.low = np.array(low)
+        self.high = np.array(high)
+        if str(self.low.dtype) not in ['float32', 'float64']:
+            self.low = self.low.astype('float32')
+            self.high = self.high.astype('float32')
+
+    def sample(self, shape):
+        shape = tuple(shape) + (self.low + self.high).shape
+        return self.low + (
+            np.random.uniform(size=shape) * (self.high - self.low)
+        )
+
+    def log_prob(self, value):
+        lb = np.less(self.low, value).astype(self.low.dtype)
+        ub = np.less(value, self.high).astype(self.low.dtype)
+        return np.log(lb * ub) - np.log(self.high - self.low)
+
+    def probs(self, value):
+        lb = np.less(self.low, value).astype(self.low.dtype)
+        ub = np.less(value, self.high).astype(self.low.dtype)
+        return (lb * ub) / (self.high - self.low)
+
+    def entropy(self):
+        return np.log(self.high - self.low)
+
+
+class UniformTest(unittest.TestCase):
+    def setUp(self, use_gpu=False, batch_size=5, dims=6):
+        self.use_gpu = use_gpu
+        if not use_gpu:
+            self.place = base.CPUPlace()
+            self.gpu_id = -1
+        else:
+            self.place = base.CUDAPlace(0)
+            self.gpu_id = 0
+
+        self.init_numpy_data(batch_size, dims)
+
+        paddle.disable_static(self.place)
+        self.init_dynamic_data(batch_size, dims)
+
+        paddle.enable_static()
+        self.test_program = base.Program()
+        self.executor = base.Executor(self.place)
+        self.init_static_data(batch_size, dims)
+
+    def init_numpy_data(self, batch_size, dims):
+        # low ans high are 'float'
+        self.low_np = np.random.uniform(-2, 1)
+        self.high_np = np.random.uniform(2, 4)
+        self.values_np = np.array([1.0]).astype('float32')
+
+    def init_dynamic_data(self, batch_size, dims):
+        self.dynamic_low = self.low_np
+        self.dynamic_high = self.high_np
+        self.dynamic_values = paddle.to_tensor(self.values_np)
+
+    def init_static_data(self, batch_size, dims):
+        self.static_low = self.low_np
+        self.static_high = self.high_np
+        with base.program_guard(self.test_program):
+            self.static_values = paddle.static.data(
+                name='values', shape=[-1], dtype='float32'
+            )
+
+    def compare_with_numpy(self, fetch_list, sample_shape=7, tolerance=1e-6):
+        sample, entropy, log_prob, probs = fetch_list
+
+        np_uniform = UniformNumpy(self.low_np, self.high_np)
+        np_sample = np_uniform.sample([sample_shape])
+        np_entropy = np_uniform.entropy()
+        np_lp = np_uniform.log_prob(self.values_np)
+        np_p = np_uniform.probs(self.values_np)
+
+        np.testing.assert_equal(sample.shape, np_sample.shape)
+        np.testing.assert_allclose(
+            entropy, np_entropy, rtol=tolerance, atol=tolerance
+        )
+        np.testing.assert_allclose(
+            log_prob, np_lp, rtol=tolerance, atol=tolerance
+        )
+        np.testing.assert_allclose(probs, np_p, rtol=tolerance, atol=tolerance)
+
+    def test_uniform_distribution_static(self, sample_shape=7, tolerance=1e-6):
+        paddle.enable_static()
+        with base.program_guard(self.test_program):
+            uniform = Uniform(self.static_low, self.static_high)
+            sample = uniform.sample([sample_shape])
+            entropy = uniform.entropy()
+            log_prob = uniform.log_prob(self.static_values)
+            probs = uniform.probs(self.static_values)
+            fetch_list = [sample, entropy, log_prob, probs]
+
+        feed_vars = {
+            'low': self.low_np,
+            'high': self.high_np,
+            'values': self.values_np,
+        }
+
+        self.executor.run(base.default_startup_program())
+        fetch_list = self.executor.run(
+            program=self.test_program, feed=feed_vars, fetch_list=fetch_list
+        )
+
+        self.compare_with_numpy(fetch_list)
+
+    def func_uniform_distribution_dygraph(self, sample_shape=7, tolerance=1e-6):
+        paddle.disable_static()
+        uniform = Uniform(self.dynamic_low, self.dynamic_high)
+        sample = uniform.sample([sample_shape]).numpy()
+        entropy = uniform.entropy().numpy()
+        log_prob = uniform.log_prob(self.dynamic_values).numpy()
+        probs = uniform.probs(self.dynamic_values).numpy()
+        fetch_list = [sample, entropy, log_prob, probs]
+
+        self.compare_with_numpy(fetch_list)
+
+    def test_uniform_distribution_dygraph(self):
+        self.setUp()
+        self.func_uniform_distribution_dygraph()
+
+
+class UniformTest2(UniformTest):
+    def init_numpy_data(self, batch_size, dims):
+        # low ans high are 'int'
+        self.low_np = int(np.random.uniform(-2, 1))
+        self.high_np = int(np.random.uniform(2, 4))
+        self.values_np = np.array([1.0]).astype('float32')
+
+
+class UniformTest3(UniformTest):
+    def init_numpy_data(self, batch_size, dims):
+        # test broadcast: low is float, high is numpy.ndarray with dtype 'float32'.
+        self.low_np = np.random.uniform(-2, 1)
+        self.high_np = np.random.uniform(5.0, 15.0, (batch_size, dims)).astype(
+            'float32'
+        )
+        self.values_np = np.random.randn(batch_size, dims).astype('float32')
+
+    def init_static_data(self, batch_size, dims):
+        self.static_low = self.low_np
+        self.static_high = self.high_np
+        with base.program_guard(self.test_program):
+            self.static_values = paddle.static.data(
+                name='values', shape=[-1, dims], dtype='float32'
+            )
+
+
+class UniformTest4(UniformTest):
+    def init_numpy_data(self, batch_size, dims):
+        # low and high are numpy.ndarray with dtype 'float32'.
+        self.low_np = np.random.randn(batch_size, dims).astype('float32')
+        self.high_np = np.random.uniform(5.0, 15.0, (batch_size, dims)).astype(
+            'float32'
+        )
+        self.values_np = np.random.randn(batch_size, dims).astype('float32')
+
+    def init_static_data(self, batch_size, dims):
+        self.static_low = self.low_np
+        self.static_high = self.high_np
+        with base.program_guard(self.test_program):
+            self.static_values = paddle.static.data(
+                name='values', shape=[-1, dims], dtype='float32'
+            )
+
+
+class UniformTest5(UniformTest):
+    def init_numpy_data(self, batch_size, dims):
+        # low and high are numpy.ndarray with dtype 'float64'.
+        self.low_np = np.random.randn(batch_size, dims).astype('float64')
+        self.high_np = np.random.uniform(5.0, 15.0, (batch_size, dims)).astype(
+            'float64'
+        )
+        self.values_np = np.random.randn(batch_size, dims).astype('float64')
+
+    def init_dynamic_data(self, batch_size, dims):
+        self.dynamic_low = self.low_np
+        self.dynamic_high = self.high_np
+        self.dynamic_values = paddle.to_tensor(self.values_np, dtype='float64')
+
+    def init_static_data(self, batch_size, dims):
+        self.static_low = self.low_np
+        self.static_high = self.high_np
+        with base.program_guard(self.test_program):
+            self.static_values = paddle.static.data(
+                name='values', shape=[-1, dims], dtype='float64'
+            )
+
+
+class UniformTest6(UniformTest):
+    def init_numpy_data(self, batch_size, dims):
+        # low and high are Tensor with dtype 'VarType.FP32'.
+        self.low_np = np.random.randn(batch_size, dims).astype('float32')
+        self.high_np = np.random.uniform(5.0, 15.0, (batch_size, dims)).astype(
+            'float32'
+        )
+        self.values_np = np.random.randn(batch_size, dims).astype('float32')
+
+    def init_dynamic_data(self, batch_size, dims):
+        self.dynamic_low = paddle.to_tensor(self.low_np)
+        self.dynamic_high = paddle.to_tensor(self.high_np)
+        self.dynamic_values = paddle.to_tensor(self.values_np)
+
+    def init_static_data(self, batch_size, dims):
+        with base.program_guard(self.test_program):
+            self.static_low = paddle.static.data(
+                name='low', shape=[-1, dims], dtype='float32'
+            )
+            self.static_high = paddle.static.data(
+                name='high', shape=[-1, dims], dtype='float32'
+            )
+            self.static_values = paddle.static.data(
+                name='values', shape=[-1, dims], dtype='float32'
+            )
+
+
+class UniformTest7(UniformTest):
+    def init_numpy_data(self, batch_size, dims):
+        # low and high are Tensor with dtype 'VarType.FP64'.
+        self.low_np = np.random.randn(batch_size, dims).astype('float64')
+        self.high_np = np.random.uniform(5.0, 15.0, (batch_size, dims)).astype(
+            'float64'
+        )
+        self.values_np = np.random.randn(batch_size, dims).astype('float64')
+
+    def init_dynamic_data(self, batch_size, dims):
+        self.dynamic_low = paddle.to_tensor(self.low_np, dtype='float64')
+        self.dynamic_high = paddle.to_tensor(self.high_np, dtype='float64')
+        self.dynamic_values = paddle.to_tensor(self.values_np, dtype='float64')
+
+    def init_static_data(self, batch_size, dims):
+        with base.program_guard(self.test_program):
+            self.static_low = paddle.static.data(
+                name='low', shape=[-1, dims], dtype='float64'
+            )
+            self.static_high = paddle.static.data(
+                name='high', shape=[-1, dims], dtype='float64'
+            )
+            self.static_values = paddle.static.data(
+                name='values', shape=[-1, dims], dtype='float64'
+            )
+
+
+class UniformTest8(UniformTest):
+    def init_numpy_data(self, batch_size, dims):
+        # low and high are Tensor with dtype 'VarType.FP64'. value's dtype is 'VarType.FP32'.
+        self.low_np = np.random.randn(batch_size, dims).astype('float64')
+        self.high_np = np.random.uniform(5.0, 15.0, (batch_size, dims)).astype(
+            'float64'
+        )
+        self.values_np = np.random.randn(batch_size, dims).astype('float32')
+
+    def init_dynamic_data(self, batch_size, dims):
+        self.dynamic_low = paddle.to_tensor(self.low_np, dtype='float64')
+        self.dynamic_high = paddle.to_tensor(self.high_np, dtype='float64')
+        self.dynamic_values = paddle.to_tensor(self.values_np, dtype='float32')
+
+    def init_static_data(self, batch_size, dims):
+        with base.program_guard(self.test_program):
+            self.static_low = paddle.static.data(
+                name='low', shape=[-1, dims], dtype='float64'
+            )
+            self.static_high = paddle.static.data(
+                name='high', shape=[-1, dims], dtype='float64'
+            )
+            self.static_values = paddle.static.data(
+                name='values', shape=[-1, dims], dtype='float32'
+            )
+
+
+class UniformTest9(UniformTest):
+    def init_numpy_data(self, batch_size, dims):
+        # low and high are numpy.ndarray with dtype 'float32'.
+        # high < low.
+        self.low_np = np.random.randn(batch_size, dims).astype('float32')
+        self.high_np = np.random.uniform(
+            -10.0, -5.0, (batch_size, dims)
+        ).astype('float32')
+        self.values_np = np.random.randn(batch_size, dims).astype('float32')
+
+    def init_static_data(self, batch_size, dims):
+        self.static_low = self.low_np
+        self.static_high = self.high_np
+        with base.program_guard(self.test_program):
+            self.static_values = paddle.static.data(
+                name='values', shape=[-1, dims], dtype='float32'
+            )
+
+
+class UniformTest10(UniformTest):
+    def init_numpy_data(self, batch_size, dims):
+        # low and high are list.
+        self.low_np = (
+            np.random.randn(batch_size, dims).astype('float32').tolist()
+        )
+        self.high_np = (
+            np.random.uniform(5.0, 15.0, (batch_size, dims))
+            .astype('float32')
+            .tolist()
+        )
+        self.values_np = np.random.randn(batch_size, dims).astype('float32')
+
+    def init_static_data(self, batch_size, dims):
+        self.static_low = self.low_np
+        self.static_high = self.high_np
+        with base.program_guard(self.test_program):
+            self.static_values = paddle.static.data(
+                name='values', shape=[-1, dims], dtype='float32'
+            )
+
+
+class UniformTest11(UniformTest):
+    def init_numpy_data(self, batch_size, dims):
+        # low and high are tuple.
+        self.low_np = tuple(
+            np.random.randn(batch_size, dims).astype('float32').tolist()
+        )
+        self.high_np = tuple(
+            np.random.uniform(5.0, 15.0, (batch_size, dims))
+            .astype('float32')
+            .tolist()
+        )
+        self.values_np = np.random.randn(batch_size, dims).astype('float32')
+
+    def init_static_data(self, batch_size, dims):
+        self.static_low = self.low_np
+        self.static_high = self.high_np
+        with base.program_guard(self.test_program):
+            self.static_values = paddle.static.data(
+                name='values', shape=[-1, dims], dtype='float32'
+            )
+
+
+class UniformTestSample(unittest.TestCase):
+    def setUp(self):
+        self.init_param()
+
+    def init_param(self):
+        self.low = 3.0
+        self.high = 4.0
+
+    def test_uniform_sample(self):
+        paddle.disable_static()
+        uniform = Uniform(low=self.low, high=self.high)
+        s = uniform.sample([100])
+        self.assertTrue((s >= self.low).all())
+        self.assertTrue((s < self.high).all())
+        paddle.enable_static()
+
+
+class UniformTestSample2(UniformTestSample):
+    def init_param(self):
+        self.low = -5.0
+        self.high = 2.0
+
+
+if __name__ == '__main__':
+    unittest.main()

test/distribution/test_distribution_variable.py

@@ -0,0 +1,79 @@
+# Copyright (c) 2022 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
+#
+#     http://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 parameterize as param
+
+import paddle
+from paddle.distribution import constraint, variable
+
+paddle.seed(2022)
+
+
+@param.param_cls(
+    (param.TEST_CASE_NAME, 'is_discrete', 'event_rank', 'constraint'),
+    [('NotImplement', False, 0, constraint.Constraint())],
+)
+class TestVariable(unittest.TestCase):
+    def setUp(self):
+        self._var = variable.Variable(
+            self.is_discrete, self.event_rank, self.constraint
+        )
+
+    @param.param_func([(1,)])
+    def test_costraint(self, value):
+        with self.assertRaises(NotImplementedError):
+            self._var.constraint(value)
+
+
+@param.param_cls(
+    (param.TEST_CASE_NAME, 'base', 'rank'), [('real_base', variable.real, 10)]
+)
+class TestIndependent(unittest.TestCase):
+    def setUp(self):
+        self._var = variable.Independent(self.base, self.rank)
+
+    @param.param_func(
+        [
+            (paddle.rand([2, 3, 4]), ValueError),
+        ]
+    )
+    def test_costraint(self, value, expect):
+        with self.assertRaises(expect):
+            self._var.constraint(value)
+
+
+@param.param_cls(
+    (param.TEST_CASE_NAME, 'vars', 'axis'), [('real_base', [variable.real], 10)]
+)
+class TestStack(unittest.TestCase):
+    def setUp(self):
+        self._var = variable.Stack(self.vars, self.axis)
+
+    def test_is_discrete(self):
+        self.assertEqual(self._var.is_discrete, False)
+
+    @param.param_func(
+        [
+            (paddle.rand([2, 3, 4]), ValueError),
+        ]
+    )
+    def test_costraint(self, value, expect):
+        with self.assertRaises(expect):
+            self._var.constraint(value)
+
+
+if __name__ == '__main__':
+    unittest.main()

test/distribution/test_kl.py

@@ -0,0 +1,158 @@
+# Copyright (c) 2021 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
+#
+#     http://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 mock_data as mock
+import numpy as np
+import parameterize as param
+import scipy.special
+import scipy.stats
+from distribution import config
+
+import paddle
+from paddle.distribution import kl
+
+np.random.seed(2022)
+paddle.seed(2022)
+paddle.set_default_dtype('float64')
+
+
+@param.place(config.DEVICES)
+@param.parameterize_cls(
+    (param.TEST_CASE_NAME, 'a1', 'b1', 'a2', 'b2'),
+    [
+        (
+            'test_regular_input',
+            6.0 * np.random.random((4, 5)) + 1e-4,
+            6.0 * np.random.random((4, 5)) + 1e-4,
+            6.0 * np.random.random((4, 5)) + 1e-4,
+            6.0 * np.random.random((4, 5)) + 1e-4,
+        ),
+    ],
+)
+class TestKLBetaBeta(unittest.TestCase):
+    def setUp(self):
+        self.p = paddle.distribution.Beta(
+            paddle.to_tensor(self.a1), paddle.to_tensor(self.b1)
+        )
+        self.q = paddle.distribution.Beta(
+            paddle.to_tensor(self.a2), paddle.to_tensor(self.b2)
+        )
+
+    def test_kl_divergence(self):
+        with paddle.base.dygraph.guard(self.place):
+            np.testing.assert_allclose(
+                paddle.distribution.kl_divergence(self.p, self.q),
+                self.scipy_kl_beta_beta(self.a1, self.b1, self.a2, self.b2),
+                rtol=config.RTOL.get(str(self.a1.dtype)),
+                atol=config.ATOL.get(str(self.a1.dtype)),
+            )
+
+    def scipy_kl_beta_beta(self, a1, b1, a2, b2):
+        return (
+            scipy.special.betaln(a2, b2)
+            - scipy.special.betaln(a1, b1)
+            + (a1 - a2) * scipy.special.digamma(a1)
+            + (b1 - b2) * scipy.special.digamma(b1)
+            + (a2 - a1 + b2 - b1) * scipy.special.digamma(a1 + b1)
+        )
+
+
+@param.place(config.DEVICES)
+@param.param_cls(
+    (param.TEST_CASE_NAME, 'conc1', 'conc2'),
+    [
+        (
+            'test-regular-input',
+            np.random.random((5, 7, 8, 10)),
+            np.random.random((5, 7, 8, 10)),
+        ),
+    ],
+)
+class TestKLDirichletDirichlet(unittest.TestCase):
+    def setUp(self):
+        self.p = paddle.distribution.Dirichlet(paddle.to_tensor(self.conc1))
+        self.q = paddle.distribution.Dirichlet(paddle.to_tensor(self.conc2))
+
+    def test_kl_divergence(self):
+        with paddle.base.dygraph.guard(self.place):
+            np.testing.assert_allclose(
+                paddle.distribution.kl_divergence(self.p, self.q),
+                self.scipy_kl_diric_diric(self.conc1, self.conc2),
+                rtol=config.RTOL.get(str(self.conc1.dtype)),
+                atol=config.ATOL.get(str(self.conc1.dtype)),
+            )
+
+    def scipy_kl_diric_diric(self, conc1, conc2):
+        return (
+            scipy.special.gammaln(np.sum(conc1, -1))
+            - scipy.special.gammaln(np.sum(conc2, -1))
+            - np.sum(
+                scipy.special.gammaln(conc1) - scipy.special.gammaln(conc2), -1
+            )
+            + np.sum(
+                (conc1 - conc2)
+                * (
+                    scipy.special.digamma(conc1)
+                    - scipy.special.digamma(np.sum(conc1, -1, keepdims=True))
+                ),
+                -1,
+            )
+        )
+
+
+class DummyDistribution(paddle.distribution.Distribution):
+    pass
+
+
+@param.place(config.DEVICES)
+@param.param_cls(
+    (param.TEST_CASE_NAME, 'p', 'q'),
+    [('test-unregister', DummyDistribution(), DummyDistribution)],
+)
+class TestDispatch(unittest.TestCase):
+    def test_dispatch_with_unregister(self):
+        with self.assertRaises(NotImplementedError):
+            paddle.distribution.kl_divergence(self.p, self.q)
+
+
+@param.place(config.DEVICES)
+@param.param_cls(
+    (param.TEST_CASE_NAME, 'p', 'q'),
+    [
+        (
+            'test-diff-dist',
+            mock.Exponential(paddle.rand((100, 200, 100)) + 1.0),
+            mock.Exponential(paddle.rand((100, 200, 100)) + 2.0),
+        ),
+        (
+            'test-same-dist',
+            mock.Exponential(paddle.to_tensor([1.0])),
+            mock.Exponential(paddle.to_tensor([1.0])),
+        ),
+    ],
+)
+class TestKLExpfamilyExpFamily(unittest.TestCase):
+    def test_kl_expfamily_expfamily(self):
+        np.testing.assert_allclose(
+            paddle.distribution.kl_divergence(self.p, self.q),
+            kl._kl_expfamily_expfamily(self.p, self.q),
+            rtol=config.RTOL.get(config.DEFAULT_DTYPE),
+            atol=config.ATOL.get(config.DEFAULT_DTYPE),
+        )
+
+
+if __name__ == '__main__':
+    unittest.main()

test/distribution/test_kl_static.py

@@ -0,0 +1,221 @@
+# Copyright (c) 2021 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
+#
+#     http://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 mock_data as mock
+import numpy as np
+import parameterize as param
+import scipy.special
+import scipy.stats
+from distribution import config
+
+import paddle
+from paddle.distribution import kl
+
+np.random.seed(2022)
+paddle.seed(2022)
+paddle.enable_static()
+
+
+@param.place(config.DEVICES)
+@param.param_cls(
+    (param.TEST_CASE_NAME, 'a1', 'b1', 'a2', 'b2'),
+    [
+        (
+            'test_regular_input',
+            6.0 * np.random.random((4, 5)) + 1e-4,
+            6.0 * np.random.random((4, 5)) + 1e-4,
+            6.0 * np.random.random((4, 5)) + 1e-4,
+            6.0 * np.random.random((4, 5)) + 1e-4,
+        ),
+    ],
+)
+class TestKLBetaBeta(unittest.TestCase):
+    def setUp(self):
+        self.mp = paddle.static.Program()
+        self.sp = paddle.static.Program()
+        self.executor = paddle.static.Executor(self.place)
+
+        with paddle.static.program_guard(self.mp, self.sp):
+            a1 = paddle.static.data('a1', self.a1.shape, dtype=self.a1.dtype)
+            b1 = paddle.static.data('b1', self.b1.shape, dtype=self.b1.dtype)
+            a2 = paddle.static.data('a2', self.a2.shape, dtype=self.a2.dtype)
+            b2 = paddle.static.data('b2', self.b2.shape, dtype=self.b2.dtype)
+
+            self.p = paddle.distribution.Beta(a1, b1)
+            self.q = paddle.distribution.Beta(a2, b2)
+            self.feeds = {
+                'a1': self.a1,
+                'b1': self.b1,
+                'a2': self.a2,
+                'b2': self.b2,
+            }
+
+    def test_kl_divergence(self):
+        with paddle.static.program_guard(self.mp, self.sp):
+            out = paddle.distribution.kl_divergence(self.p, self.q)
+            self.executor.run(self.sp)
+            [out] = self.executor.run(
+                self.mp, feed=self.feeds, fetch_list=[out]
+            )
+
+            np.testing.assert_allclose(
+                out,
+                self.scipy_kl_beta_beta(self.a1, self.b1, self.a2, self.b2),
+                rtol=config.RTOL.get(str(self.a1.dtype)),
+                atol=config.ATOL.get(str(self.a1.dtype)),
+            )
+
+    def scipy_kl_beta_beta(self, a1, b1, a2, b2):
+        return (
+            scipy.special.betaln(a2, b2)
+            - scipy.special.betaln(a1, b1)
+            + (a1 - a2) * scipy.special.digamma(a1)
+            + (b1 - b2) * scipy.special.digamma(b1)
+            + (a2 - a1 + b2 - b1) * scipy.special.digamma(a1 + b1)
+        )
+
+
+@param.place(config.DEVICES)
+@param.param_cls(
+    (param.TEST_CASE_NAME, 'conc1', 'conc2'),
+    [
+        (
+            'test-regular-input',
+            np.random.random((5, 7, 8, 10)),
+            np.random.random((5, 7, 8, 10)),
+        ),
+    ],
+)
+class TestKLDirichletDirichlet(unittest.TestCase):
+    def setUp(self):
+        self.mp = paddle.static.Program()
+        self.sp = paddle.static.Program()
+        self.executor = paddle.static.Executor(self.place)
+        with paddle.static.program_guard(self.mp, self.sp):
+            conc1 = paddle.static.data(
+                'conc1', self.conc1.shape, self.conc1.dtype
+            )
+            conc2 = paddle.static.data(
+                'conc2', self.conc2.shape, self.conc2.dtype
+            )
+            self.p = paddle.distribution.Dirichlet(conc1)
+            self.q = paddle.distribution.Dirichlet(conc2)
+            self.feeds = {'conc1': self.conc1, 'conc2': self.conc2}
+
+    def test_kl_divergence(self):
+        with paddle.static.program_guard(self.mp, self.sp):
+            out = paddle.distribution.kl_divergence(self.p, self.q)
+            self.executor.run(self.sp)
+            [out] = self.executor.run(
+                self.mp, feed=self.feeds, fetch_list=[out]
+            )
+            np.testing.assert_allclose(
+                out,
+                self.scipy_kl_diric_diric(self.conc1, self.conc2),
+                rtol=config.RTOL.get(str(self.conc1.dtype)),
+                atol=config.ATOL.get(str(self.conc1.dtype)),
+            )
+
+    def scipy_kl_diric_diric(self, conc1, conc2):
+        return (
+            scipy.special.gammaln(np.sum(conc1, -1))
+            - scipy.special.gammaln(np.sum(conc2, -1))
+            - np.sum(
+                scipy.special.gammaln(conc1) - scipy.special.gammaln(conc2), -1
+            )
+            + np.sum(
+                (conc1 - conc2)
+                * (
+                    scipy.special.digamma(conc1)
+                    - scipy.special.digamma(np.sum(conc1, -1, keepdims=True))
+                ),
+                -1,
+            )
+        )
+
+
+class DummyDistribution(paddle.distribution.Distribution):
+    pass
+
+
+@param.place(config.DEVICES)
+@param.param_cls((param.TEST_CASE_NAME, 'p', 'q'), ['test-dispatch-exception'])
+class TestDispatch(unittest.TestCase):
+    def setUp(self):
+        self.mp = paddle.static.Program()
+        self.sp = paddle.static.Program()
+        self.executor = paddle.static.Executor(self.place)
+        with paddle.static.program_guard(self.mp, self.sp):
+            self.p = DummyDistribution()
+            self.q = DummyDistribution()
+
+    def test_dispatch_with_unregister(self):
+        with (
+            self.assertRaises(NotImplementedError),
+            paddle.static.program_guard(self.mp, self.sp),
+        ):
+            out = paddle.distribution.kl_divergence(self.p, self.q)
+            self.executor.run(self.sp)
+            self.executor.run(self.mp, feed={}, fetch_list=[out])
+
+
+@param.place(config.DEVICES)
+@param.param_cls(
+    (config.TEST_CASE_NAME, 'rate1', 'rate2'),
+    [
+        (
+            'test-diff-dist',
+            np.random.rand(100, 200, 100) + 1.0,
+            np.random.rand(100, 200, 100) + 2.0,
+        ),
+        ('test-same-dist', np.array([1.0]), np.array([1.0])),
+    ],
+)
+class TestKLExpfamilyExpFamily(unittest.TestCase):
+    def setUp(self):
+        self.mp = paddle.static.Program()
+        self.sp = paddle.static.Program()
+        self.executor = paddle.static.Executor(self.place)
+        with paddle.static.program_guard(self.mp, self.sp):
+            rate1 = paddle.static.data(
+                'rate1', shape=self.rate1.shape, dtype=self.rate1.dtype
+            )
+            rate2 = paddle.static.data(
+                'rate2', shape=self.rate2.shape, dtype=self.rate2.dtype
+            )
+            self.p = mock.Exponential(rate1)
+            self.q = mock.Exponential(rate2)
+            self.feeds = {'rate1': self.rate1, 'rate2': self.rate2}
+
+    def test_kl_expfamily_expfamily(self):
+        with paddle.static.program_guard(self.mp, self.sp):
+            out1 = paddle.distribution.kl_divergence(self.p, self.q)
+            out2 = kl._kl_expfamily_expfamily(self.p, self.q)
+            self.executor.run(self.sp)
+            [out1, out2] = self.executor.run(
+                self.mp, feed=self.feeds, fetch_list=[out1, out2]
+            )
+
+            np.testing.assert_allclose(
+                out1,
+                out2,
+                rtol=config.RTOL.get(config.DEFAULT_DTYPE),
+                atol=config.ATOL.get(config.DEFAULT_DTYPE),
+            )
+
+
+if __name__ == '__main__':
+    unittest.main()