Add distributions of Categorical and MultivariateNormal (#18263)

* add_distributions_of_normal_and_uniform

* paddle/fluid/API.spec

* modify API.spec

* modified paddle/fluid/API.spec, test=develop

* modify paddle/fluid/API.spec, test=develop

* modify paddle/fluid/API.spec, test=develop

* fix some comment, test=develop

* modify API.spec, test=develop

* Add distributions of Categorical and MultivariateNormal test=develop

* fix pylint codestyle test=develop

* fix conflict file test=develop

* edit API.spec test=develop

* improve sample code test=develop

* modify api.spec test=develop

paddle/fluid/API.spec

@@ -439,6 +439,18 @@ paddle.fluid.layers.Normal.entropy (ArgSpec(args=['self'], varargs=None, keyword
 paddle.fluid.layers.Normal.kl_divergence (ArgSpec(args=['self', 'other'], varargs=None, keywords=None, defaults=None), ('document', '2e8845cdf1129647e6fa6e816876cd3b'))
 paddle.fluid.layers.Normal.log_prob (ArgSpec(args=['self', 'value'], varargs=None, keywords=None, defaults=None), ('document', 'b79091014ceaffb6a7372a198a341c23'))
 paddle.fluid.layers.Normal.sample (ArgSpec(args=['self', 'shape', 'seed'], varargs=None, keywords=None, defaults=(0,)), ('document', 'adac334af13f6984e991b3ecf12b8cb7'))
+paddle.fluid.layers.Categorical ('paddle.fluid.layers.distributions.Categorical', ('document', '865c9dac8af6190e05588486ba091ee8'))
+paddle.fluid.layers.Categorical.__init__ (ArgSpec(args=['self', 'logits'], varargs=None, keywords=None, defaults=None), ('document', '933b96c9ebab8e2c1f6007a50287311e'))
+paddle.fluid.layers.Categorical.entropy (ArgSpec(args=['self'], varargs=None, keywords=None, defaults=None), ('document', 'b360a2a7a4da07c2d268b329e09c82c1'))
+paddle.fluid.layers.Categorical.kl_divergence (ArgSpec(args=['self', 'other'], varargs=None, keywords=None, defaults=None), ('document', 'c2c4c37376584178025f0a4a61c4b862'))
+paddle.fluid.layers.Categorical.log_prob (ArgSpec(args=['self', 'value'], varargs=None, keywords=None, defaults=None), ('document', 'c0edd2e2fc76711477b32dc4da9de768'))
+paddle.fluid.layers.Categorical.sample (ArgSpec(args=['self'], varargs=None, keywords=None, defaults=None), ('document', '08a2bbcaa20ee176ee7ec3d05737a0f6'))
+paddle.fluid.layers.MultivariateNormalDiag ('paddle.fluid.layers.distributions.MultivariateNormalDiag', ('document', 'f6ee0e8b2898796dcff2a68c9fda19f0'))
+paddle.fluid.layers.MultivariateNormalDiag.__init__ (ArgSpec(args=['self', 'loc', 'scale'], varargs=None, keywords=None, defaults=None), ('document', '6adf97f83acf6453d4a6a4b1070f3754'))
+paddle.fluid.layers.MultivariateNormalDiag.entropy (ArgSpec(args=['self'], varargs=None, keywords=None, defaults=None), ('document', '3c679b573ba975c5067c8ebfd4354b02'))
+paddle.fluid.layers.MultivariateNormalDiag.kl_divergence (ArgSpec(args=['self', 'other'], varargs=None, keywords=None, defaults=None), ('document', 'd9190d29dbd54c81f747a6436c35f062'))
+paddle.fluid.layers.MultivariateNormalDiag.log_prob (ArgSpec(args=['self', 'value'], varargs=None, keywords=None, defaults=None), ('document', 'c0edd2e2fc76711477b32dc4da9de768'))
+paddle.fluid.layers.MultivariateNormalDiag.sample (ArgSpec(args=['self'], varargs=None, keywords=None, defaults=None), ('document', '08a2bbcaa20ee176ee7ec3d05737a0f6'))
 paddle.fluid.contrib.InitState ('paddle.fluid.contrib.decoder.beam_search_decoder.InitState', ('document', '3afd1f84232718e628e9e566941c5f05'))
 paddle.fluid.contrib.InitState.__init__ (ArgSpec(args=['self', 'init', 'shape', 'value', 'init_boot', 'need_reorder', 'dtype'], varargs=None, keywords=None, defaults=(None, None, 0.0, None, False, 'float32')), ('document', '6adf97f83acf6453d4a6a4b1070f3754'))
 paddle.fluid.contrib.StateCell ('paddle.fluid.contrib.decoder.beam_search_decoder.StateCell', ('document', 'ecd0066c02867d445d7b461e28220c50'))

python/paddle/fluid/layers/distributions.py

@@ -22,7 +22,7 @@ import math
 import numpy as np
 import warnings
 
-__all__ = ['Uniform', 'Normal']
+__all__ = ['Uniform', 'Normal', 'Categorical', 'MultivariateNormalDiag']
 
 
 class Distribution(object):
@@ -396,3 +396,208 @@ class Normal(Distribution):
         t1 = (self.loc - other.loc) / other.scale
         t1 = (t1 * t1)
         return 0.5 * (var_ratio + t1 - 1. - nn.log(var_ratio))
+
+
+class Categorical(Distribution):
+    """
+    Categorical distribution is a discrete probability distribution that 
+    describes the possible results of a random variable that can take on 
+    one of K possible categories, with the probability of each category 
+    separately specified.
+
+    Args:
+        logits(list|numpy.ndarray|Variable): The logits input of categorical distribution.
+
+    Examples:
+        .. code-block:: python
+
+          import numpy as np
+          from paddle.fluid import layers
+          from paddle.fluid.layers import Categorical
+
+          a_logits_npdata = np.array([-0.602,-0.602], dtype="float32")
+          a_logits_tensor = layers.create_tensor(dtype="float32")
+          layers.assign(a_logits_npdata, a_logits_tensor)
+
+          b_logits_npdata = np.array([-0.102,-0.112], dtype="float32")
+          b_logits_tensor = layers.create_tensor(dtype="float32")
+          layers.assign(b_logits_npdata, b_logits_tensor)
+          
+          a = Categorical(a_logits_tensor)
+          b = Categorical(b_logits_tensor)
+
+          a.entropy()
+          # [0.6931472] with shape: [1]
+
+          b.entropy()
+          # [0.6931347] with shape: [1]
+
+          a.kl_divergence(b)
+          # [1.2516975e-05] with shape: [1]
+
+    """
+
+    def __init__(self, logits):
+        """
+        Args:
+            logits: A float32 tensor
+        """
+        if self._validate_args(logits):
+            self.logits = logits
+        else:
+            self.logits = self._to_variable(logits)[0]
+
+    def kl_divergence(self, other):
+        """The KL-divergence between two Categorical distributions.
+
+        Args:
+            other (Categorical): instance of Categorical.
+
+        Returns:
+            Variable: kl-divergence between two Categorical distributions.
+
+        """
+        assert isinstance(other, Categorical)
+
+        logits = self.logits - nn.reduce_max(self.logits, dim=-1, keep_dim=True)
+        other_logits = other.logits - nn.reduce_max(
+            other.logits, dim=-1, keep_dim=True)
+        e_logits = ops.exp(logits)
+        other_e_logits = ops.exp(other_logits)
+        z = nn.reduce_sum(e_logits, dim=-1, keep_dim=True)
+        other_z = nn.reduce_sum(other_e_logits, dim=-1, keep_dim=True)
+        prob = e_logits / z
+        kl = nn.reduce_sum(
+            prob * (logits - nn.log(z) - other_logits + nn.log(other_z)),
+            dim=-1,
+            keep_dim=True)
+
+        return kl
+
+    def entropy(self):
+        """Shannon entropy in nats.
+
+        Returns:
+          Variable: Shannon entropy of Categorical distribution.
+
+        """
+        logits = self.logits - nn.reduce_max(self.logits, dim=-1, keep_dim=True)
+        e_logits = ops.exp(logits)
+        z = nn.reduce_sum(e_logits, dim=-1, keep_dim=True)
+        prob = e_logits / z
+        entropy = -1.0 * nn.reduce_sum(
+            prob * (logits - nn.log(z)), dim=-1, keep_dim=True)
+
+        return entropy
+
+
+class MultivariateNormalDiag(Distribution):
+    """
+    A multivariate normal (also called Gaussian) distribution parameterized by a mean vector
+    and a covariance matrix.
+
+    Args:
+        loc(list|numpy.ndarray|Variable): The mean of multivariateNormal distribution.
+        scale(list|numpy.ndarray|Variable): The positive definite diagonal covariance matrix of
+        multivariateNormal distribution.
+
+    Examples:
+        .. code-block:: python
+    
+            import numpy as np
+            from paddle.fluid import layers
+            from paddle.fluid.layers import MultivariateNormalDiag
+
+            a_loc_npdata = np.array([0.3,0.5],dtype="float32")
+            a_loc_tensor = layers.create_tensor(dtype="float32")
+            layers.assign(a_loc_npdata, a_loc_tensor)
+
+
+            a_scale_npdata = np.array([[0.4,0],[0,0.5]],dtype="float32")
+            a_scale_tensor = layers.create_tensor(dtype="float32")
+            layers.assign(a_scale_npdata, a_scale_tensor)
+
+            b_loc_npdata = np.array([0.2,0.4],dtype="float32")
+            b_loc_tensor = layers.create_tensor(dtype="float32")
+            layers.assign(b_loc_npdata, b_loc_tensor)
+
+            b_scale_npdata = np.array([[0.3,0],[0,0.4]],dtype="float32")
+            b_scale_tensor = layers.create_tensor(dtype="float32")
+            layers.assign(b_scale_npdata, b_scale_tensor)
+
+            a = MultivariateNormalDiag(a_loc_tensor, a_scale_tensor)
+            b = MultivariateNormalDiag(b_loc_tensor, b_scale_tensor)
+            
+            a.entropy()
+            # [2.033158] with shape: [1]
+            b.entropy()
+            # [1.7777451] with shaoe: [1]
+
+            a.kl_divergence(b)
+            # [0.06542051] with shape: [1]
+       
+    """
+
+    def __init__(self, loc, scale):
+        if self._validate_args(loc, scale):
+            self.loc = loc
+            self.scale = scale
+        else:
+            self.loc, self.scale = self._to_variable(loc, scale)
+
+    def _det(self, value):
+
+        batch_shape = list(value.shape)
+        one_all = tensor.ones(shape=batch_shape, dtype=self.loc.dtype)
+        one_diag = tensor.diag(
+            tensor.ones(
+                shape=[batch_shape[0]], dtype=self.loc.dtype))
+        det_diag = nn.reduce_prod(value + one_all - one_diag)
+
+        return det_diag
+
+    def _inv(self, value):
+
+        batch_shape = list(value.shape)
+        one_all = tensor.ones(shape=batch_shape, dtype=self.loc.dtype)
+        one_diag = tensor.diag(
+            tensor.ones(
+                shape=[batch_shape[0]], dtype=self.loc.dtype))
+        inv_diag = nn.elementwise_pow(value, (one_all - 2 * one_diag))
+
+        return inv_diag
+
+    def entropy(self):
+        """Shannon entropy in nats.
+
+        Returns:
+          Variable: Shannon entropy of Multivariate Normal distribution.
+
+        """
+        entropy = 0.5 * (
+            self.scale.shape[0] *
+            (1.0 + math.log(2 * math.pi)) + nn.log(self._det(self.scale)))
+
+        return entropy
+
+    def kl_divergence(self, other):
+        """The KL-divergence between two Multivariate Normal distributions.
+
+        Args:
+            other (MultivariateNormalDiag): instance of Multivariate Normal.
+
+        Returns:
+            Variable: kl-divergence between two Multivariate Normal distributions.
+
+        """
+        assert isinstance(other, MultivariateNormalDiag)
+
+        tr_cov_matmul = nn.reduce_sum(self._inv(other.scale) * self.scale)
+        loc_matmul_cov = nn.matmul((other.loc - self.loc),
+                                   self._inv(other.scale))
+        tri_matmul = nn.matmul(loc_matmul_cov, (other.loc - self.loc))
+        k = list(self.scale.shape)[0]
+        ln_cov = nn.log(self._det(other.scale)) - nn.log(self._det(self.scale))
+        kl = 0.5 * (tr_cov_matmul + tri_matmul - k + ln_cov)
+
+        return kl

python/paddle/fluid/tests/unittests/test_distributions.py

@@ -88,6 +88,68 @@ class NormalNumpy(DistributionNumpy):
         return 0.5 * (var_ratio + t1 - 1 - np.log(var_ratio))
 
 
+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. * np.sum(prob * (logits - np.log(z)), axis=-1, keepdims=True)
+
+    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 MultivariateNormalDiagNumpy(DistributionNumpy):
+    def __init__(self, loc, scale):
+        self.loc = np.array(loc).astype('float32')
+        self.scale = np.array(scale).astype('float32')
+
+    def _det(self, value):
+        batch_shape = list(value.shape)
+        one_all = np.ones(shape=batch_shape, dtype='float32')
+        one_diag = np.eye(batch_shape[0], dtype='float32')
+        det_diag = np.prod(value + one_all - one_diag)
+
+        return det_diag
+
+    def _inv(self, value):
+        batch_shape = list(value.shape)
+        one_all = np.ones(shape=batch_shape, dtype='float32')
+        one_diag = np.eye(batch_shape[0], dtype='float32')
+        inv_diag = np.power(value, (one_all - 2 * one_diag))
+
+        return inv_diag
+
+    def entropy(self):
+        return 0.5 * (self.scale.shape[0] *
+                      (1.0 + np.log(np.array(2 * math.pi).astype('float32'))
+                       ) + np.log(self._det(self.scale)))
+
+    def kl_divergence(self, other):
+        tr_cov_matmul = np.sum(self._inv(other.scale) * self.scale)
+        loc_matmul_cov = np.matmul((other.loc - self.loc),
+                                   self._inv(other.scale))
+        tri_matmul = np.matmul(loc_matmul_cov, (other.loc - self.loc))
+        k = list(self.scale.shape)[0]
+        ln_cov = np.log(self._det(other.scale)) - np.log(self._det(self.scale))
+        kl = 0.5 * (tr_cov_matmul + tri_matmul - k + ln_cov)
+
+        return kl
+
+
 class DistributionTest(unittest.TestCase):
     def setUp(self, use_gpu=False):
         self.use_gpu = use_gpu
@@ -410,6 +472,105 @@ class DistributionTest(unittest.TestCase):
         np.testing.assert_allclose(
             output_lp_variable, gt_lp, rtol=tolerance, atol=tolerance)
 
+    def test_categorical_distribution(self,
+                                      batch_size=2,
+                                      dims=3,
+                                      tolerance=1e-6):
+        test_program = fluid.Program()
+
+        logits_np = np.random.randn(batch_size, dims).astype('float32')
+        other_logits_np = np.random.randn(batch_size, dims).astype('float32')
+
+        with fluid.program_guard(test_program):
+            logits = layers.data(name='logits', shape=[dims], dtype='float32')
+            other_logits = layers.data(
+                name='other_logits', shape=[dims], dtype='float32')
+
+            categorical_np = Categorical(logits_np)
+            other_categorical_np = Categorical(other_logits_np)
+
+            entropy_np = categorical_np.entropy()
+            kl_np = categorical_np.kl_divergence(other_categorical_np)
+
+        self.executor.run(fluid.default_main_program())
+
+        np_categorical = CategoricalNumpy(logits_np)
+        np_other_categorical = CategoricalNumpy(other_logits_np)
+        gt_entropy_np = np_categorical.entropy()
+        gt_kl_np = np_categorical.kl_divergence(np_other_categorical)
+
+        # result calculated by paddle
+        [output_entropy_np,
+         output_kl_np] = self.executor.run(program=test_program,
+                                           feed={'logits': logits_np},
+                                           fetch_list=[entropy_np, kl_np])
+        np.testing.assert_allclose(
+            output_entropy_np, gt_entropy_np, rtol=tolerance)
+        np.testing.assert_allclose(output_kl_np, gt_kl_np, rtol=tolerance)
+
+    def test_multivariateNormalDiag_distribution(self,
+                                                 batch_size=2,
+                                                 tolerance=1e-6):
+        test_program = fluid.Program()
+
+        loc_np = np.random.random(batch_size, ).astype('float32')
+        scale_np = np.diag(np.random.random(batch_size, )).astype('float32')
+        other_loc_np = np.random.random(batch_size, ).astype('float32')
+        other_scale_np = np.diag(np.random.random(batch_size, )).astype(
+            'float32')
+
+        with fluid.program_guard(test_program):
+            loc = layers.data(
+                name='loc',
+                shape=[batch_size, ],
+                dtype='float32',
+                append_batch_size=False)
+            scale = layers.data(
+                name='scale',
+                shape=[batch_size, batch_size],
+                dtype='float32',
+                append_batch_size=False)
+            other_loc = layers.data(
+                name='other_loc',
+                shape=[batch_size, ],
+                dtype='float32',
+                append_batch_size=False)
+            other_scale = layers.data(
+                name='other_scale',
+                shape=[batch_size, batch_size],
+                dtype='float32',
+                append_batch_size=False)
+
+            multivariate_np = MultivariateNormalDiag(loc, scale)
+            other_multivariate_np = MultivariateNormalDiag(other_loc,
+                                                           other_scale)
+
+            entropy_np = multivariate_np.entropy()
+            other_entropy_np = other_multivariate_np.entropy()
+            kl_np = multivariate_np.kl_divergence(other_multivariate_np)
+
+        self.executor.run(fluid.default_main_program())
+
+        np_multivariate = MultivariateNormalDiagNumpy(loc_np, scale_np)
+        np_other_multivariate = MultivariateNormalDiagNumpy(other_loc_np,
+                                                            other_scale_np)
+        gt_entropy_np = np_multivariate.entropy()
+        gt_kl_np = np_multivariate.kl_divergence(np_other_multivariate)
+
+        # result calculated by paddle
+        [output_entropy_np,
+         output_kl_np] = self.executor.run(program=test_program,
+                                           feed={
+                                               'loc': loc_np,
+                                               'scale': scale_np,
+                                               'other_loc': other_loc_np,
+                                               'other_scale': other_scale_np
+                                           },
+                                           fetch_list=[entropy_np, kl_np])
+        np.testing.assert_allclose(
+            output_entropy_np, gt_entropy_np, rtol=tolerance)
+        np.testing.assert_allclose(output_kl_np, gt_kl_np, rtol=tolerance)
+
 
 if __name__ == '__main__':
     unittest.main()