PaddlePaddle · luotao1 · Mar 29, 2023 · Mar 29, 2023 · Mar 29, 2023
diff --git a/rfcs/APIs/20230318_api_design_for_i0&i0e.md b/rfcs/APIs/20230318_api_design_for_i0&i0e.md
@@ -1,12 +1,12 @@
 # paddle.i0 和 paddle.i0e 设计文档
 
-| API 名称      |         paddle.i0 / paddle.i0e           |
-| ------------ | ---------------------------------------- |
-| 提交作者    | PommesPeter                               |
-| 提交时间    | 2023-03-18                                |
-| 版本号      | V1.0                                      |
-| 依赖飞桨版本 | develop                                   |
-| 文件名      | 20230318_api_design_for_i0&i0e.md      |
+| API 名称     | paddle.i0 / paddle.i0e            |
+| ------------ | --------------------------------- |
+| 提交作者     | PommesPeter                       |
+| 提交时间     | 2023-03-29                        |
+| 版本号       | V1.1                              |
+| 依赖飞桨版本 | develop                           |
+| 文件名       | 20230318_api_design_for_i0&i0e.md |
 
 # 一、概述
 
@@ -16,7 +16,7 @@
 
 ## 2、功能目标
 
-根据输入的 tensor，计算其每个元素的第一类零阶修正贝塞尔函数（对应api：`paddle.i0`）和第一类指数缩放的零阶修正贝塞尔函数（对应api：`paddle.i0e`）。
+根据输入的 tensor，计算其每个元素的第一类零阶修正贝塞尔函数（对应 api：`paddle.i0`）和第一类指数缩放的零阶修正贝塞尔函数（对应 api：`paddle.i0e`）。
 
 ## 3、意义
 
@@ -30,19 +30,19 @@
 
 ## PyTorch
 
-PyTorch 中有 `torch.special.i0` 的API，详细参数为 `torch.special.i0(input, *, out=None) → Tensor`。
-PyTorch 中有 `torch.special.i0e` 的API，详细参数为 `torch.special.i0e(input, *, out=None) → Tensor`。
+PyTorch 中有 `torch.special.i0` 的 API，详细参数为 `torch.special.i0(input, *, out=None) → Tensor`。
+PyTorch 中有 `torch.special.i0e` 的 API，详细参数为 `torch.special.i0e(input, *, out=None) → Tensor`。
 
 ### i0
 
 > Computes the zeroth order modified Bessel function of the first kind for each element of `input`.
-在实现方法上，PyTorch 是通过 C++ API 组合实现的
 
+在实现方法上，PyTorch 是通过 C++ API 组合实现的
 实现代码：
 
 [代码位置](https://github.com/pytorch/pytorch/blob/HEAD/torch/csrc/api/include/torch/special.h#L456-L471)
 
-```cpp
+````cpp
 /// Computes the zeroth order modified Bessel function of the first kind of
 /// input, elementwise See
 /// https://pytorch.org/docs/master/special.html#torch.special.i0
@@ -59,7 +59,7 @@ inline Tensor i0(const Tensor& self) {
 inline Tensor& i0_out(Tensor& result, const Tensor& self) {
   return torch::special_i0_out(result, self);
 }
-```
+````
 
 [代码位置](https://github.com/pytorch/pytorch/blob/HEAD/aten/src/ATen/native/Math.h#L1458-L1474)
 
@@ -87,7 +87,7 @@ calc_i0(T _x) {
 
 [代码位置](https://github.com/pytorch/pytorch/blob/HEAD/torch/csrc/api/include/torch/special.h#L490-L505)
 
-```cpp
+````cpp
 /// Computes the exponentially scaled zeroth order modified Bessel function of
 /// the first kind See
 /// https://pytorch.org/docs/master/special.html#torch.special.i0e.
@@ -104,7 +104,7 @@ inline Tensor i0e(const Tensor& self) {
 inline Tensor& i0e_out(Tensor& result, const Tensor& self) {
   return torch::special_i0e_out(result, self);
 }
-```
+````
 
 [代码位置](https://github.com/pytorch/pytorch/blob/HEAD/aten/src/ATen/native/Math.h#L101-L125)
 
@@ -135,6 +135,7 @@ if (x <= T{8.0}) {
     return chbevl(y, coefficients, int{30});
 }
 ```
+
 参数表：
 
 - input (Tensor) – the input tensor
@@ -228,7 +229,7 @@ double x;
 
     if (x < 0)
         x = -x;
-    if (x <= 8.0) 
+    if (x <= 8.0)
     {
         y = (x / 2.0) - 2.0;
         return (exp(x) * chbevl(y, A, 30));
@@ -250,24 +251,54 @@ double x;
 }
 ```
 
+## TensorFlow
+
+### i0
+
+TensorFlow 中已有 `i0` 算子的反向传播实现，可作为参考。
+
+```python
+@ops.RegisterGradient("BesselI0")
+def _BesselI0Grad(op, grad):
+  """Compute gradient of bessel_i0(x) with respect to its argument."""
+  x = op.inputs[0]
+  with ops.control_dependencies([grad]):
+    partial_x = special_math_ops.bessel_i1(x)
+    return grad * partial_x
+```
+
+### i0e
+
+TensorFlow 中已有 `i0e` 算子的反向传播实现，可作为参考。
+
+```python
+@ops.RegisterGradient("BesselI0e")
+def _BesselI0eGrad(op, grad):
+  """Compute gradient of bessel_i0e(x) with respect to its argument."""
+  x = op.inputs[0]
+  y = op.outputs[0]
+  with ops.control_dependencies([grad]):
+    partial_x = (special_math_ops.bessel_i1e(x) - math_ops.sign(x) * y)
+    return grad * partial_x
+```
+
 # 四、对比分析
 
 ## 共同点
 
-- 都能通过输入的 tensor，计算其每个元素的第一类零阶修正贝塞尔函数和第一类指数缩放的零阶修正贝塞尔函数
-- 都支持对向量或者矩阵进行运算，对向量或矩阵中的每一个元素计算，是一个 element-wise 的操作
-- 都有提供对 Python 的调用接口
+- 都能通过输入的 tensor，计算其每个元素的第一类零阶修正贝塞尔函数和第一类指数缩放的零阶修正贝塞尔函数。
+- 都支持对向量或者矩阵进行运算，对向量或矩阵中的每一个元素计算，是一个 element-wise 的操作。
+- 都有提供对 Python 的调用接口。
 
 ## 不同点
 
 - PyTorch 是在 C++ API 基础上实现，使用 Python 调用 C++ 对应的接口。
-- Scipy 则是参考数学函数库 Cephes 实现对应的 API
+- Scipy 则是参考数学函数库 Cephes 实现对应的 API。
 
 # 五、设计思路与实现方案
 
 ## 命名与参数设计
 
-
 添加 Python API
 
 ```python
@@ -284,42 +315,65 @@ paddle.i0e(
 )
 ```
 
-## 底层OP设计
+## 底层 OP 设计
+
+考虑到 `paddle` 本身已经引入 Eigen3 库，并且 Eigen3 库已经包含相应的实现逻辑，可参考 Eigen3 库进行 OP 设计。
 
 ### i0
-参考 Scipy 的 i0 和 i0e 进行实现，实现位置为 Paddle repo 的 `paddle/phi/kernels`。可参考 Cephes 数学函数库的实现方式。
 
 观察公式：
 
 $$out_i=I_0(x)=\sum_{k=0}^{\infin}\frac{(x_i^2/4)^k}{(k!)^2}$$
 
-可以发现右侧可以使用切比雪夫级数进行计算，故提前计算切比雪夫级数并代入公式即可。
+对于 i0，基于实现的 i0e 乘以缩放系数即可。公式为：
 
-### i0e
+参考 Eigen3 的 i0 函数进行实现，需要调用 `i0e` 算子，实现位置为 Paddle repo 的 `paddle/phi/kernels`。前向传播计算过程为：
+
+1. 对输入计算绝对值 `y = abs(x)`
+2. 对 `y` 计算自然指数 `y' = exp(y)`
+3. 调用 `i0e` 得到结果 `out = y' * i0e(y)`
 
-对于 i0e，基于实现的 i0 乘以缩放系数即可。公式为：
+对 `i0(x)` 函数求偏导为 `i1(x)`, 故反向传播计算过程为：
+
+1. 调用 `i1` 计算 `i0` 的偏导数 `partial_x`
+2. 计算 `grad * partial_x`
+
+### i0e
 
 $$out_i=\exp(-|x|) * I_0(x_i)=\exp(-|x|) * \sum_{k=0}^{\infin}\frac{(x_i^2/4)^k}{(k!)^2}$$
 
+将定义域分为 $[0, 8]$ 和 $[8, \infty]$ 两个区间，在每个区间内部分别通过 Chebyshev 多项式展开计算系数，故提前计算切比雪夫多项式展开系数并代入公式即可。
+
+参考 Eigen3 的 i0 函数进行实现，需要调用 `i0e` 算子，实现位置为 Paddle repo 的 `paddle/phi/kernels`。前向传播计算过程为：
 
-## API实现方案
+1. 对输入计算绝对值 `y = abs(x)`
+2. 如果输入的数值小于 8，则使用 $[0, 8]$ 区间上的 Chebyshev 多项式系数计算 `chbevl((y / 2.0) - 2.0, A)`
+3. 如果输入的数值大于 8，则使用 $[8, \infty]$ 区间上的 Chebyshev 多项式系数计算 `chbevl((32.0 / y) - 2.0, B) / sqrt(y)`
 
-参考 Scipy 进行实现，该 API 实现于 `python/paddle/tensor/math.py` 和 `paddle/phi/kernels`。
+对 `i0e(x)` 函数求偏导为 `i1(x) - sign(x) * y`, 故反向传播计算过程为：
 
+1. 调用 `i1` 计算 `i0` 的偏导数 `partial_x`
+2. 计算 `grad * partial_x`
+
+## API 实现方案
+
+参考 Eigen3 进行实现，该 API 实现于 `python/paddle/tensor/math.py` 和 `paddle/phi/kernels`。
 
 # 六、测试和验收的考量
 
 测试需要考虑的 case 如下：
 
 - 输出数值结果的一致性和数据类型是否正确，使用 scipy 作为参考标准
-- 参数 `x` 的数据类型准确性判断
+- 对不同 dtype 的输入数据 `x` 进行计算精度检验 (float32, float64)
+- 对不同范围内的输入数据进行计算精度检验 ($[0, 8]$, $[8, \infty]$)
 - 输入输出的容错性与错误提示信息
-- 输出Dtype错误或不兼容时抛出异常
+- 输出 Dtype 错误或不兼容时抛出异常
 - 保证调用属性时是可以被正常找到的
+- 覆盖静态图和动态图测试场景
 
 # 七、可行性分析和排期规划
 
-方案主要依赖现有 Scipy 代码参考实现。工期上可以满足在当前版本周期内开发完成。
+方案主要依赖现有 Eigen3 库代码参考实现。工期上可以满足在当前版本周期内开发完成。
 
 # 八、影响面
 
@@ -335,7 +389,8 @@ $$out_i=\exp(-|x|) * I_0(x_i)=\exp(-|x|) * \sum_{k=0}^{\infin}\frac{(x_i^2/4)^k}
 [torch.special.i0e](https://pytorch.org/docs/stable/special.html#torch.special.i0e)
 
 [scipy.special.i0](https://docs.scipy.org/doc/scipy/reference/generated/scipy.special.i0.html#scipy-special-i0)
-
 [scipy.special.i0e](https://docs.scipy.org/doc/scipy/reference/generated/scipy.special.i0e.html#scipy.special.i0e)
 
-[paddle.complex](https://github.com/PaddlePaddle/Paddle/blob/develop/python/paddle/tensor/creation.py#L2160-L2209)
+[TensorFlow API 文档](https://www.tensorflow.org/api_docs/python/tf/math/bessel_i1)
+
+[Rational Approximations for the Modified Bessel Function of the First Kind – I1(x) for Computations with Double Precision](https://www.bing.com/search?q=%22Rational+Approximations+for+the+Modified+Bessel+Function+of+the+First+Kind+%0D%0A%23+++++-+I1%28x%29+for+Computations+with+Double+Precision%22+by+Pavel+Holoborodko&qs=n&form=QBRE&sp=-1&lq=1&pq=%22rational+approximations+for+the+modified+bessel+function+of+the+first+kind+%23+-+i1%28x%29+for+computations+with+double+precision%22+by+pavel+holoborodko&sc=0-146&sk=&cvid=2507FDB609B24E39BEA10DE5D3482BD8&ghsh=0&ghacc=0&ghpl=)