为paddle新增LogNormal概率分布

在paddle已有的概率分布的基础上，开发LogNormal概率分

AI Studio

342人浏览 · 2022-12-05 13:26:24

AI Studio · 2022-12-05 13:26:24 发布

1 任务介绍

在 paddle 中，paddle.distribution 目录下包含了随机变量的概率分布、随机变量的变换、KL 散度相关 API。本次任务的目的是在现有的概率分布方案的基础上，实现 Log Normal 概率分布。

任务要求我们熟悉python，了解概率分布的基本知识。

任务的难度一般，即使我们不了解深度学习的相关知识，通过学习和模仿已有的概率分布方案，也能够完成本任务。

2 设计文档

2.1 初步设计

任务计划新增 LogNormal API，用于 Log Normal 分布的概率统计与随机采样，包括如下方法：

mean 计算均值
variance 计算方差
sample 随机采样
rsample 重参数化采
prob 概率密度
log_prob 对数概率密度
entropy 熵计算
kl_divergence 相对熵计算

这些方法可以在paddle.distribution 的已有概率分布中看到，新的概率分布的代码风格及设计思路与已有概率分布保持一致。

为了更好地完成开发任务，我们需要先阅读贡献指南，了解代码贡献的流程，API的设计规范和文档的书写规范。

2.2 框架学习

paddle 的概率分布方案较多，为了明确初步的学习范围，我们需要先了解什么是Log Normal 分布。

Normal 分布是大家熟知的概率分布，它的定义如下，

若随机变量 $X$ 服从均值为 $\mu$ ，方差为 $\sigma^2$ 的 Normal 分布，则 $X$ 的概率密度函数为
$f(x)=\frac{1}{\sqrt{2\pi}\sigma}e^{-\frac{(x-\mu)^2}{2\sigma^{2}}}$

记作 $\sim N(\mu, \sigma^2)$ 。

Log Normal分布的基础就是 Normal 分布：

如果随机变量 $X$ 满足 $lnX\sim N(\mu, \sigma^2)$ ，则称 $X$ 服从 Log Normal 分布， $X$ 的概率密度函数为

$\frac{1}{\sigma x \sqrt{2\pi}} e^{(-\frac{(ln(x)-\mu)^2}{2\sigma^2})}$

记作 $\sim N(\mu, \sigma^2)$ 。

可见 Normal 分布是 Log Normal 分布的基础，我们可先学习 Normal 概率分布的实现。

paddle.distribution中已经实现了 Normal 概率分布，可以通过源代码中的例子学习 Normal API 的使用。

import paddle
from paddle.distribution import Normal
# Define a single scalar Normal distribution.
dist = Normal(loc=0., scale=3.)
# Define a batch of two scalar valued Normals.
# The first has mean 1 and standard deviation 11, the second 2 and 22.
dist = Normal(loc=[1., 2.], scale=[11., 22.])
# Get 3 samples, returning a 3 x 2 tensor.
dist.sample([3])
# Define a batch of two scalar valued Normals.
# Both have mean 1, but different standard deviations.
dist = Normal(loc=1., scale=[11., 22.])
# Complete example
value_tensor = paddle.to_tensor([0.8], dtype="float32")
normal_a = Normal([0.], [1.])
normal_b = Normal([0.5], [2.])
sample = normal_a.sample([2])
# a random tensor created by normal distribution with shape: [2, 1]
entropy = normal_a.entropy()
# [1.4189385] with shape: [1]
lp = normal_a.log_prob(value_tensor)
# [-1.2389386] with shape: [1]
p = normal_a.probs(value_tensor)
# [0.28969154] with shape: [1]
kl = normal_a.kl_divergence(normal_b)
# [0.34939718] with shape: [1]

Normal 类的源代码

需要注意Normal paddle 的旧版本1.8的 API，paddle 现在已经发布2.0版本，我们开发 LogNormal 要使用 paddle 2.0 的 API。阅读升级指南了解 paddle 新旧版本的差异信息，大致浏览一遍新旧 API 的映射表。

paddle.distribution 中的 Multinomial 类是基于 paddle 2.0 ，可以作为参考。

Distribution 类是 paddle 概率分布的基类， Normal 类继承了 Distribution 类。Normal 类中实现了 prob，log_prob，entropy，kl_divergence等方法。

2.3 竞品分析

在设计 paddle 新的 API 之前，应该分析和学习其它深度学习框架同类 API 的设计方案，对比不同方案的优势与不足，从而总结出自己的方案。

深度学习框架 pytorch 和 tenserflow 中有概率分布 API，我们可以阅读它们的 LogNormal API 文档。

Tensorflow 的 LogNormal 文档

Pytorch 的 LogNormal 文档

两个框架的 LogNormal 类的初始化参数类似，都需要传入 loc 和 scale 两个参数。

#tensorflow
tfp.distributions.LogNormal(
    loc,
    scale,
    validate_args=False,
    allow_nan_stats=True,
    name='LogNormal'
)

#pytorch
torch.distributions.log_normal.LogNormal(
    loc, 
    scale, 
    validate_args=None
)

两者的LogNormal 类的基类都是TransformedDistribution，从文档中可了解到，如果概率分布 $P (y)$ 可以表示成一个基础概率分布 $P (x)$ ，以及一个可微可逆的变换 $Y = g (X)$ ，那么可以用通过继承 TransformedDistribution类，将概率分布 $P (x)$ 变换为新的概率分布 $P (y)$ 。

LogNormal 是满足使用TransformedDistribution 类的条件。

$\sim Normal(\mu, \sigma)$
$\sim LogNormal(\mu, \sigma)$

再看 pytorch 的 LogNormal 的初始化方法，代码中将 Normal 作为基础分布，并用 ExpTransform 方法指示了变换过程。

    def __init__(self, loc, scale, validate_args=None):
        base_dist = Normal(
            loc,
            scale,
            validate_args=validate_args
        )
        super(LogNormal, self).__init__(
            base_dist,
            ExpTransform(),
            validate_args=validate_args
        )

弄清楚了 pytorch 和 tensorflow 的 LogNormal 的实现方案，我们回到 paddle，明确目前 paddle 的 API 支持状况。

在paddle.distribution 中已经实现了TransformedDistribution 类，并且已经定义了 ExpTransform 变换。

TransformedDistribution的源代码

ExpTransform的源代码

从TransformedDistribution 的代码中可以看到，我们继承了该类，大部分方法就不需要自己实现了，类中已经定义了基础分布和新的分布对应方法之间的变换。

paddle 中暂时还没有TransformedDistribution 的子类可供参考，但我们可以阅读 TransformedDistribution 类的测试代码，学习如何使用该类。

TransformedDistribution的测试代码

@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
        )
    ...

paddle 的 TransformedDistribution 的使用方法与 pytroch 类似，并且也测试参数恰好是实现 Log Normal 分布所需的参数，我们就大体明确了 LogNormal 的实现方法。

2.4 实现方案

参考 paddle 的 Normal 的实现，我们可以初步设计出 API 的名称和参数。

paddle.distribution.LogNormal(loc, scale)

参数 loc, scale 分别为基础分布 Normal 的均值和标准差。

例如，随机变量 $X$ 服从 Log Normal 分布，即 $\sim N(\mu, \sigma^2)$ ，对应的参数 $loc=\mu$ ， $scale=\sigma$ 。

我们通过继承TransformedDistribution 类实现 LogNormal 类，将 Normal 分布变换为 Log Normal 分布。

class LogNormal(TransformedDistribution):
  def __init__(self, loc, scale):
    base_dist = Normal(loc, scale)
    super(LogNormal, self).__init__(
        base_dist,
        [paddle.distribution.ExpTransform()]
    )
    
    ...

接下来我们逐个设计LogNormal 类的方法。

LogNormal 类的初始化参数有 $l o c$ 和 $s c a l e$ ，记参数 $loc=\mu$ ， $scale=\sigma$ 。

mean 计算均值

均值的计算方法： $e^ {\mu + \frac{\sigma^2}{2}}$

variance 计算方差

方差的计算方法： $e^{2\mu + \sigma^2}(e^{\sigma^2} - 1)$

entropy 熵计算

熵的计算方法：基础分布 Normal 的熵与 $\mu$ 求和。

kl_divergence 相对熵计算

与对应的基础分布Normal的计算逻辑相同，复用即可。

sample 随机采样

继承父类 TransformedDistribution 的 sample 方法，将基础分布 Normal 的 sample 结果进行变换，得到 Log Normal 分布的 sample 结果。

rsample 重参数化采样
继承父类
prob 概率密度
继承父类
log_prob 对数概率密度
继承父类

2.5 单测方案

在框架开发中，对新增 API 的详细测试是必不可少的。
我们通过参考Multinomial 类的测试代码，以及 tensorflow 的 LogNormal 类的测试代码，设计出 paddle 的LogNormal 类的测试方案。

Multinomial的测试代码（动态图）

Multinomial的测试代码（静态图）

Tensorflow 的 LogNormal 测试代码

参考已有概率分布方案的测试代码，LogNormal 类测试以 Numpy 作为基准，验证API的正确性。测试的主要内容如下，

使用 Numpy 实现所有 Log Normal 的API，集成为 LogNormalNumpy 类，用以验证本次任务开发的 API 的正确性。
使用同样的参数实例化 LogNormal 类和 LogNormalNumpy 类，并调用 mean、variance、entropy、prob、log_prob、kl_divergence方法，测试结果是否一致。
使用 LogNormal 类的 sample 方法生成样本，测试这些这样的均值和标准差是否正确，并且进行KS检验。

单测使用parameterize包进行参数化测试，大致结构如下。

from parameterize import TEST_CASE_NAME, parameterize_cls, place, xrand

@place(config.DEVICES)
@parameterize_cls(
    (TEST_CASE_NAME, 'loc', 'scale', 'value'),
    [
        ... #测试的参数
        
    ],
)
class LogNormalTest(unittest.TestCase):
    def setUp(self):
        ...

    def test_mean(self):
        ...

    def test_variance(self):
        ...

    def test_entropy(self):
        ...

    def test_probs(self):
        ...

    def test_log_prob(self):
        ...

3 代码开发

3.1 API开发

我个人的开发环境是 Ubuntu20 + VS Code，使用的是 docker 环境的 padlde，配合 VS Code 的 docker 插件。

依照设计文档，我们就可以开始代码开发了。

LogNormal 类开发较为简单，实现一些前文提到的公式，就能完成mean，variance，entropy，probs方法的开发。log_prob，sample， rsample方法通过继承父类实现。kl_divergence方法复用 Normal 分布的即可。

import paddle
from paddle.distribution.normal import Normal
from paddle.distribution.transform import ExpTransform
from paddle.distribution.transformed_distribution import TransformedDistribution


class LogNormal(TransformedDistribution):

    def __init__(self, loc, scale):
        self._base = Normal(loc=loc, scale=scale)
        self.loc = self._base.loc
        self.scale = self._base.scale
        super().__init__(self._base, [ExpTransform()])

    @property
    def mean(self):
        """Mean of lognormal distribuion.
        """
        return paddle.exp(self._base.mean + self._base.variance / 2)

    @property
    def variance(self):
        """Variance of lognormal distribution.
        """
        return paddle.expm1(self._base.variance) * paddle.exp(
            2 * self._base.mean + self._base.variance
        )

    def entropy(self):
        r"""Shannon entropy in nats.
        """
        return self._base.entropy() + self._base.mean

    def probs(self, value):
        """Probability density/mass function.
        """
        return paddle.exp(self.log_prob(value))

    def kl_divergence(self, other):
        r"""The KL-divergence between two lognormal distributions.
        """
        return self._base.kl_divergence(other._base)

代码开发完成了，我们还要按照贡献指南的注释规范，补充和修改注释，书写中文和英文文档。

3.2 单测开发

我们以 Numpy 作为基准，验证API的正确性，使用 Numpy 实现 LogNormal 类的方法。

import numpy as np
from test_distribution import DistributionNumpy

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))

LogNormal 类的单测要分别在动态图模式和静态图模型下进行。

首先进行动态图测试，测试mean、variance、entropy、prob、log_prob、kl_divergence方法。依照其它概率分布的单测代码，使用 assertEqual 和 np.testing.assert_allclose 进行数据比对。

import math
import unittest

import config
import numpy as np
import scipy.stats
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


@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.fluid.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.fluid.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)),
            )

测试sample 和 rsample 方法。我们首先测试生成的样本的均值和方差是否符合预期，然后对样本进行KS检验。KS检验是较为关键的一步，我们使用 scipy.stats 模块进行检验。

@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 = 100000
        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

测试 kl_divergence 方法，由于 Normal 和 LogNormal 的kl_divergence 定义相同，我们可以参考Normal 的测试代码。

@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))

到此动态图模式的单测就完成了，静态图也是类似的。

import unittest

import config
import numpy as np
import scipy.stats
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))),
    ],
)
class TestLogNormal(unittest.TestCase):
    def setUp(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 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))],
)
class TestLogNormalSample(unittest.TestCase):
    def setUp(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 = 100000
            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 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)),
        ),
    ],
)
class TestLogNormalKL(unittest.TestCase):
    def setUp(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 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))

单测代码的覆盖率应尽量做到 100%。

4 成果展示

至此，我们完成了 LogNormal API 的开发任务，完成项目的编译后，我们可以运行一遍 LogNormal API 的所有方法，以及单测代码。

import paddle
from paddle.distribution import LogNormal
# Define a single scalar LogNormal distribution.
dist = LogNormal(loc=0., scale=3.)
# Define a batch of two scalar valued LogNormals.
# The underlying Normal of first has mean 1 and standard deviation 11, the underlying Normal of second 2 and 22.
dist = LogNormal(loc=[1., 2.], scale=[11., 22.])
# Get 3 samples, returning a 3 x 2 tensor.
dist.sample((3, ))
# Define a batch of two scalar valued LogNormals.
# Their underlying Normal have mean 1, but different standard deviations.
dist = LogNormal(loc=1., scale=[11., 22.])
# Complete example
value_tensor = paddle.to_tensor([0.8], dtype="float32")
lognormal_a = LogNormal([0.], [1.])
lognormal_b = LogNormal([0.5], [2.])
sample = lognormal_a.sample((2, ))
# a random tensor created by lognormal distribution with shape: [2, 1]
entropy = lognormal_a.entropy()
# [1.4189385] with shape: [1]
lp = lognormal_a.log_prob(value_tensor)
# [-0.72069150] with shape: [1]
p = lognormal_a.probs(value_tensor)
# [0.48641577] with shape: [1]
kl = lognormal_a.kl_divergence(lognormal_b)
# [0.34939718] with shape: [1]