Repository: secretflow/spu (language: python)

Issue:
使用SPU实现chi-squared统计量
> 此 ISSUE 为 [隐语开源共建计划（SecretFlow Open Source Contribution Plan，简称 SF OSCP）](https://github.com/orgs/secretflow/discussions/1181)任务 ISSUE，欢迎社区开发者参与共建～
> 若有感兴趣想要认领的任务，但还未报名，辛苦先完成[报名](https://www.wjx.top/vm/QKio0dq.aspx# )进行哈～


## 任务介绍
- 任务名称：使用SPU实现chi-squared统计量
- 技术方向：SPU/SML
- 任务难度：进阶🌟🌟
- 任务完成时间：4周


## 详细要求
- 安全性（尽量少 reveal）
- 功能性：同时返回chi2统计量和对应的p value
- 正确性：与sklearn的结果数值上基本一致（在可接受误差范围内）
- 代码规范：Python 代码需要使用 black+isort 进行格式化（流水线包含代码规范检查卡点）
- 提交说明：关联该 isuue 并提交代码至 https://github.com/secretflow/spu/tree/main/sml(具体提交文件夹可与reviewer讨论）
- 特殊说明：若某个特性有特殊的限制，如需要 FM128，需要更多 fxp 等需要在注释文档中明确说明


## 能力要求
- 熟悉经典的统计学理论
- 熟悉 JAX 或 NumPy，可以使用 NumPy 实现算法

## 操作说明
- 可参考sklearn中的说明: https://scikit-learn.org/stable/modules/classes.html#module-sklearn.feature_selection
- 操作指引：https://www.secretflow.org.cn/docs/spu/latest/en-US/getting_started/tutorials/develop_your_first_mpc_application
- 范例：https://github.com/secretflow/spu/blob/main/sml/linear_model/simple_sgd.py
- 可参考隐语mooc2期视频: 
  - https://www.bilibili.com/video/BV1ba4y1S73z/?spm_id_from=333.788
  - https://www.bilibili.com/video/BV1Qe411D7MG/?spm_id_from=333.788
- 贡献指引文档: https://github.com/secretflow/spu/blob/main/sml/development.md


Discussion / hints:
lizzy-0323 Give it to me
Please tell me which folder should I submit to?
> Please tell me which folder should I submit to?

hello, you can first create a new folder named "feature_selection", with a new file named "univariate_selection.py", and then implement the chi2 function in this file.

In other words, the chi2 should be implemented in **sml/feature_selection/univariate_selection.py**;

(BTW, don't forget the essential tests/emulations folders and the bazel files.)

Best wishes
Thank you!

# 关于docker环境上无法运行编译运行的问题
## os
Ubuntu 20.04.6 LTS
## 运行命令
`bazel run -c opt //sml/cluster/tests:kmeans_test`
## 报错
<img width="1448" alt="image" src="https://github.com/secretflow/spu/assets/58327725/6a8df31d-98f1-45fc-93dc-434474b9df8c">
ps:前置库均已安装完毕，貌似是网络问题？但没有找到解决方法0.0
> # 关于docker环境上无法运行编译运行的问题
> ## os
> Ubuntu 20.04.6 LTS
> 
> ## 运行命令
> `bazel run -c opt //sml/cluster/tests:kmeans_test`
> 
> ## 报错
> <img alt="image" width="1448" src="https://private-user-images.githubusercontent.com/58327725/318171490-6a8df31d-98f1-45fc-93dc-434474b9df8c.png?jwt=eyJhbGciOiJIUzI1NiIsInR5cCI6IkpXVCJ9.eyJpc3MiOiJnaXRodWIuY29tIiwiYXVkIjoicmF3LmdpdGh1YnVzZXJjb250ZW50LmNvbSIsImtleSI6ImtleTUiLCJleHAiOjE3MTE3ODE0NDAsIm5iZiI6MTcxMTc4MTE0MCwicGF0aCI6Ii81ODMyNzcyNS8zMTgxNzE0OTAtNmE4ZGYzMWQtOThmMS00NWZjLTkzZGMtNDM0NDc0YjlkZjhjLnBuZz9YLUFtei1BbGdvcml0aG09QVdTNC1ITUFDLVNIQTI1NiZYLUFtei1DcmVkZW50aWFsPUFLSUFWQ09EWUxTQTUzUFFLNFpBJTJGMjAyNDAzMzAlMkZ1cy1lYXN0LTElMkZzMyUyRmF3czRfcmVxdWVzdCZYLUFtei1EYXRlPTIwMjQwMzMwVDA2NDU0MFomWC1BbXotRXhwaXJlcz0zMDAmWC1BbXotU2lnbmF0dXJlPWIxZmUxMjYwMTNjNjNlNmQ4NzBlYzNlYzM5NjExNDY0ZjFlYTNlMzg2OTRkZjE4M2M5NWQ0ODMxMzU1YjNiODAmWC1BbXotU2lnbmVkSGVhZGVycz1ob3N0JmFjdG9yX2lkPTAma2V5X2lkPTAmcmVwb19pZD0wIn0.Ve-Xo2XbEUcELPIj7HZztEP6h9tuEdfQV9PeW5MgyiQ"> ps:前置库均已安装完毕，貌似是网络问题？但没有找到解决方法0.0

请确保 docker 可以正常访问 github
docker是可以正常访问github的，clone等操作也没问题，并且使用了代理。
> docker是可以正常访问github的，clone等操作也没问题，并且使用了代理。

hmmm, 看起来和上周末 xz 发生的故事有点关系，我来看看
@lizzy-0323 main 上最新的代码再试一下？
> @lizzy-0323 main 上最新的代码再试一下？

好的 已经成功运行起来啦！
请问如果我要统计标签的种类，是否在调用chi2函数前要求传入num_class这一参数呢，因为需要满足jit的条件，似乎大部分的方法都不可以使用。
> 请问如果我要统计标签的种类，是否在调用chi2函数前要求传入num_class这一参数呢，因为需要满足jit的条件，似乎大部分的方法都不可以使用。

是的，，不过建议是手动指定y的具体取值范围（比如传一个list进来），可以作为一个额外参数传入哈～
抱歉我没太懂您的意思，比如y的取值有0,1,3三种时，传入的是[0,2]?还是说把y所有可能的取值去重之后传入呢？
这是我目前的实现，但似乎无法通过spu环境：
![image](https://github.com/secretflow/spu/assets/58327725/6dace19d-00ba-422f-9ffb-5113c5cc2079)

> 抱歉我没太懂您的意思，比如y的取值有0,1,3三种时，传入的是[0,2]?还是说把y所有可能的取值去重之后传入呢？ 这是我目前的实现，但似乎无法通过spu环境： ![image](https://private-user-images.githubusercontent.com/58327725/320707308-6dace19d-00ba-422f-9ffb-5113c5cc2079.png?jwt=eyJhbGciOiJIUzI1NiIsInR5cCI6IkpXVCJ9.eyJpc3MiOiJnaXRodWIuY29tIiwiYXVkIjoicmF3LmdpdGh1YnVzZXJjb250ZW50LmNvbSIsImtleSI6ImtleTUiLCJleHAiOjE3MTI2MzM3NjUsIm5iZiI6MTcxMjYzMzQ2NSwicGF0aCI6Ii81ODMyNzcyNS8zMjA3MDczMDgtNmRhY2UxOWQtMDBiYS00MjJmLTlmZmItNTExM2M1Y2MyMDc5LnBuZz9YLUFtei1BbGdvcml0aG09QVdTNC1ITUFDLVNIQTI1NiZYLUFtei1DcmVkZW50aWFsPUFLSUFWQ09EWUxTQTUzUFFLNFpBJTJGMjAyNDA0MDklMkZ1cy1lYXN0LTElMkZzMyUyRmF3czRfcmVxdWVzdCZYLUFtei1EYXRlPTIwMjQwNDA5VDAzMzEwNVomWC1BbXotRXhwaXJlcz0zMDAmWC1BbXotU2lnbmF0dXJlPWUzMTlhNDE5YTAyMDA3ZjEyMjYxNGRhZGZhNDUxODU4ZTMwY2IzNGZlYzI3YTQ2NGM5N2Y4MmVjMjc1OWY1MjgmWC1BbXotU2lnbmVkSGVhZGVycz1ob3N0JmFjdG9yX2lkPTAma2V5X2lkPTAmcmVwb19pZD0wIn0.OknwrG6aivHAFqw6PXX03LelNTmH6QlPlq0XisF-evQ)

y的取值有0,1,3： 我的意思是直接传入[0,1,3]；
BTW，您可以先在函数上加@jax.jit，只有明文下也是jitable的函数才有可能在SPU下正确执行。
需要对chi2函数也加上@jax.jit吗？确保其中所有的函数都是jitable的？
这样的话，好像jnp.zeros等函数也无法使用
您在测试明文实现阶段可以加jit，spu在运行的时候自动会jit，所以最终您不需要在代码里显式的加这个。

您在编程的时候要注意所有的shape应该是确定的，只要不涉及dynamic shape，jnp.zeros等函数是能使用的
## 目前进展
算法明文测试通过，和sklearn结果基本一致，以下是两种方法得到的结果：
![image](https://github.com/secretflow/spu/assets/58327725/9e4b06c4-5f07-4946-a936-ca200a3457f7)
## 目前遇到的问题
dot函数和outer函数均不能使用，查阅文档后发现spu是支持dot函数的，outer函数没有查到，目前的代码如下所示：
```python
def chi2(X, y, label_lst):
    total_samples = X.shape[0]
    num_feature = X.shape[1]
    num_class = len(label_lst)
    X = jnp.array(X)
    y = jnp.array(y)
    feature_count = jnp.sum(X, axis=0)
    class_prob = jnp.mean(y, axis=0)
    # observed = jnp.zeros((num_feature, num_class))
    observed = jnp.dot(X.T, y)
    expected = jnp.zeros((num_feature, num_class))
    # expected = jnp.outer(feature_count, class_prob)
    # print(observed)
    # print(expected)
    chi2_stats = (observed - expected) ** 2 / expected
    chi2_stats = jnp.nansum(chi2_stats, axis=1)
    df = num_class - 1
    p_value = sf(chi2_stats, df=df)
    return chi2_stats, p_value
```
函数参数：y数组传入的是一个one-hot编码后的数组，并还需要传入一个label数组作为辅助
其中，为了确认两个函数不能使用，在测试时选择对另一个数组采用zeros初始化的方法。
想请教一下问题大概出在哪里
> ## 目前进展
> 算法明文测试通过，和sklearn结果基本一致，以下是两种方法得到的结果： ![image](https://private-user-images.githubusercontent.com/58327725/320768856-9e4b06c4-5f07-4946-a936-ca200a3457f7.png?jwt=eyJhbGciOiJIUzI1NiIsInR5cCI6IkpXVCJ9.eyJpc3MiOiJnaXRodWIuY29tIiwiYXVkIjoicmF3LmdpdGh1YnVzZXJjb250ZW50LmNvbSIsImtleSI6ImtleTUiLCJleHAiOjE3MTI3MTUxODYsIm5iZiI6MTcxMjcxNDg4NiwicGF0aCI6Ii81ODMyNzcyNS8zMjA3Njg4NTYtOWU0YjA2YzQtNWYwNy00OTQ2LWE5MzYtY2EyMDBhMzQ1N2Y3LnBuZz9YLUFtei1BbGdvcml0aG09QVdTNC1ITUFDLVNIQTI1NiZYLUFtei1DcmVkZW50aWFsPUFLSUFWQ09EWUxTQTUzUFFLNFpBJTJGMjAyNDA0MTAlMkZ1cy1lYXN0LTElMkZzMyUyRmF3czRfcmVxdWVzdCZYLUFtei1EYXRlPTIwMjQwNDEwVDAyMDgwNlomWC1BbXotRXhwaXJlcz0zMDAmWC1BbXotU2lnbmF0dXJlPTBkNzI4ZDZkNTMwMmQ1NGNjZDdiOGY0ZjI4NjM4ODNiNDc3MWQyYTJiNTUwZTI4OWRjZTdhODQ5MmFkNDdjZDkmWC1BbXotU2lnbmVkSGVhZGVycz1ob3N0JmFjdG9yX2lkPTAma2V5X2lkPTAmcmVwb19pZD0wIn0.ChVseo9NgSkQRqQUZqVpomkWkFunamWUyPWyszkjqEA)
> 
> ## 目前遇到的问题
> dot函数和outer函数均不能使用，查阅文档后发现spu是支持dot函数的，outer函数没有查到，目前的代码如下所示：
> 
> ```python
> def chi2(X, y, label_lst):
>     total_samples = X.shape[0]
>     num_feature = X.shape[1]
>     num_class = len(label_lst)
>     X = jnp.array(X)
>     y = jnp.array(y)
>     feature_count = jnp.sum(X, axis=0)
>     class_prob = jnp.mean(y, axis=0)
>     # observed = jnp.zeros((num_feature, num_class))
>     observed = jnp.dot(X.T, y)
>     expected = jnp.zeros((num_feature, num_class))
>     # expected = jnp.outer(feature_count, class_prob)
>     # print(observed)
>     # print(expected)
>     chi2_stats = (observed - expected) ** 2 / expected
>     chi2_stats = jnp.nansum(chi2_stats, axis=1)
>     df = num_class - 1
>     p_value = sf(chi2_stats, df=df)
>     return chi2_stats, p_value
> ```
> 
> 函数参数：y数组传入的是一个one-hot编码后的数组，并还需要传入一个label数组作为辅助 其中，为了确认两个函数不能使用，在测试时选择对另一个数组采用zeros初始化的方法。 想请教一下问题大概出在哪里

hello，抱歉昨天有事没有及时回复；

首先，dot和outer函数都是支持的，您要不直接发起一个PR，把完整的代码推上来，否则只从这个代码片段找不太出来原因-。-
> > ## 目前进展
> > 算法明文测试通过，和sklearn结果基本一致，以下是两种方法得到的结果： ![image](https://private-user-images.githubusercontent.com/58327725/320768856-9e4b06c4-5f07-4946-a936-ca200a3457f7.png?jwt=eyJhbGciOiJIUzI1NiIsInR5cCI6IkpXVCJ9.eyJpc3MiOiJnaXRodWIuY29tIiwiYXVkIjoicmF3LmdpdGh1YnVzZXJjb250ZW50LmNvbSIsImtleSI6ImtleTUiLCJleHAiOjE3MTI3MTUxODYsIm5iZiI6MTcxMjcxNDg4NiwicGF0aCI6Ii81ODMyNzcyNS8zMjA3Njg4NTYtOWU0YjA2YzQtNWYwNy00OTQ2LWE5MzYtY2EyMDBhMzQ1N2Y3LnBuZz9YLUFtei1BbGdvcml0aG09QVdTNC1ITUFDLVNIQTI1NiZYLUFtei1DcmVkZW50aWFsPUFLSUFWQ09EWUxTQTUzUFFLNFpBJTJGMjAyNDA0MTAlMkZ1cy1lYXN0LTElMkZzMyUyRmF3czRfcmVxdWVzdCZYLUFtei1EYXRlPTIwMjQwNDEwVDAyMDgwNlomWC1BbXotRXhwaXJlcz0zMDAmWC1BbXotU2lnbmF0dXJlPTBkNzI4ZDZkNTMwMmQ1NGNjZDdiOGY0ZjI4NjM4ODNiNDc3MWQyYTJiNTUwZTI4OWRjZTdhODQ5MmFkNDdjZDkmWC1BbXotU2lnbmVkSGVhZGVycz1ob3N0JmFjdG9yX2lkPTAma2V5X2lkPTAmcmVwb19pZD0wIn0.ChVseo9NgSkQRqQUZqVpomkWkFunamWUyPWyszkjqEA)
> > ## 目前遇到的问题
> > dot函数和outer函数均不能使用，查阅文档后发现spu是支持dot函数的，outer函数没有查到，目前的代码如下所示：
> > ```python
> > def chi2(X, y, label_lst):
> >     total_samples = X.shape[0]
> >     num_feature = X.shape[1]
> >     num_class = len(label_lst)
> >     X = jnp.array(X)
> >     y = jnp.array(y)
> >     feature_count = jnp.sum(X, axis=0)
> >     class_prob = jnp.mean(y, axis=0)
> >     # observed = jnp.zeros((num_feature, num_class))
> >     observed = jnp.dot(X.T, y)
> >     expected = jnp.zeros((num_feature, num_class))
> >     # expected = jnp.outer(feature_count, class_prob)
> >     # print(observed)
> >     # print(expected)
> >     chi2_stats = (observed - expected) ** 2 / expected
> >     chi2_stats = jnp.nansum(chi2_stats, axis=1)
> >     df = num_class - 1
> >     p_value = sf(chi2_stats, df=df)
> >     return chi2_stats, p_value
> > ```
> > 
> > 
> >     
> >       
> >     
> > 
> >       
> >     
> > 
> >     
> >   
> > 函数参数：y数组传入的是一个one-hot编码后的数组，并还需要传入一个label数组作为辅助 其中，为了确认两个函数不能使用，在测试时选择对另一个数组采用zeros初始化的方法。 想请教一下问题大概出在哪里
> 
> hello，抱歉昨天有事没有及时回复；
> 
> 首先，dot和outer函数都是支持的，您要不直接发起一个PR，把完整的代码推上来，否则只从这个代码片段找不太出来原因-。-


好的，pr已提交，实在麻烦您！

经沟通，确认由[lizzy-0323](https://github.com/lizzy-0323)持续进行改项目的开发，任务时间放宽至5月中旬，感谢对隐语的关注和支持～
目前的实现如下，按照scipy源码写了一个python版本
```python
import jax
import jax.numpy as jnp
import numpy as np
from sklearn.datasets import load_iris

MACHEP = 0.0000001  # the machine roundoff error / tolerance
BIG = 4.503599627370496e15
BIGINV = 2.22044604925031308085e-16


def _sf(x, df):
    def body_func(x_i):
        if x_i < 0:
            return 1
        if x_i == 0:
            return 0
        if df <= 0:
            raise ValueError("Domain error")
        if x_i < 1.0 or x_i < df:
            return 1.0 - _igam(0.5 * df, 0.5 * x_i)
        return _igamc(df * 0.5, x_i * 0.5)

    return jax.vmap(body_func)(x)


def _igam(a, x):
    # Compute  x**a * exp(-x) / Gamma(a)
    ax = jax.exp(a * jax.log(x) - x - jax.lgamma(a))

    # Power series
    r = a
    c = 1.0
    ans = 1.0

    while True:
        r += 1.0
        c *= x / r
        ans += c
        if c / ans <= MACHEP:
            return ans * ax / a


def _igamc(a, x):
    # Compute  x**a * exp(-x) / Gamma(a) ,ax
    ax = jnp.exp(a * jnp.log(x) - x - jax.lax.lgamma(a))
    # print(ax)
    # Continued fraction
    y = 1.0 - a
    z = x + y + 1.0
    c = 0.0
    pkm2 = 1.0
    qkm2 = x
    pkm1 = x + 1.0
    qkm1 = z * x
    ans = pkm1 / qkm1
    while True:
        c += 1.0
        y += 1.0
        z += 2.0
        yc = y * c
        pk = pkm1 * z - pkm2 * yc
        qk = qkm1 * z - qkm2 * yc
        if qk != 0:
            r = pk / qk
            t = abs((ans - r) / r)
            ans = r
        else:
            t = 1.0
        pkm2 = pkm1
        pkm1 = pk
        qkm2 = qkm1
        qkm1 = qk
        if abs(pk) > BIG:
            pkm2 *= BIGINV
            pkm1 *= BIGINV
            qkm2 *= BIGINV
            qkm1 *= BIGINV
        if t <= MACHEP:
            return ans * ax


def chi2(X, y, label_lst):
    """
    Calculate the chi-squared statistic and p-value for feature independence testing.

    Parameters
    ----------
    X : array-like, shape (n_samples, n_features)
        The feature matrix from which to calculate the chi-squared statistic. Each row
        represents a sample and each column a feature.

    y : array-like, shape (n_samples,)
        The class labels of each sample as integers.

    label_lst : list of str, length = n_classes
        A list of class labels corresponding to the columns of y.

    Returns
    -------
    chi2_stats : array, shape (n_features,)
        The chi-squared statistic for each feature, indicating the degree of
        association between the feature and the class labels.

    p_value : array, shape (n_features,)
        The p-value associated with each chi-squared statistic, which can be used
        to test the null hypothesis that the features are independent of the
        class labels.
    """
    total_samples = X.shape[0]
    num_feature = X.shape[1]
    num_class = len(label_lst)
    # one hot encoding
    y = jnp.eye(num_class)[y]
    X = jnp.array(X)
    y = jnp.array(y)
    feature_count = jnp.sum(X, axis=0)
    class_prob = jnp.mean(y, axis=0)
    # Calculate the observed frequency count
    observed = jnp.dot(X.T, y)
    expected = jnp.outer(feature_count, class_prob)
    # Calculate the chi-squared statistic
    chi2_stats = (observed - expected) ** 2 / expected
    # Sum over class dimensions
    chi2_stats = jnp.nansum(chi2_stats, axis=1)
    # Degrees of freedom
    df = num_class - 1
    # Calculate the p-value for each feature
    p_value = _sf(chi2_stats, df=df)
    return chi2_stats, p_value


if __name__ == '__main__':
    x, y = load_iris(return_X_y=True)
    label_lst = np.unique(y)
    chi2_stats, p_value = chi2(x, y, label_lst)
    print(chi2_stats, p_value)
    from sklearn.feature_selection import chi2 as chi2_sklearn

    sklearn_chi2_stats, sklearn_p_value = chi2_sklearn(x, y)
    print(sklearn_chi2_stats, sklearn_p_value)
```
> 目前的实现如下，按照scipy源码写了一个python版本
> 
> ```python
> import jax
> import jax.numpy as jnp
> import numpy as np
> from sklearn.datasets import load_iris
> 
> MACHEP = 0.0000001  # the machine roundoff error / tolerance
> BIG = 4.503599627370496e15
> BIGINV = 2.22044604925031308085e-16
> 
> 
> def _sf(x, df):
>     def body_func(x_i):
>         if x_i < 0:
>             return 1
>         if x_i == 0:
>             return 0
>         if df <= 0:
>             raise ValueError("Domain error")
>         if x_i < 1.0 or x_i < df:
>             return 1.0 - _igam(0.5 * df, 0.5 * x_i)
>         return _igamc(df * 0.5, x_i * 0.5)
> 
>     return jax.vmap(body_func)(x)
> 
> 
> def _igam(a, x):
>     # Compute  x**a * exp(-x) / Gamma(a)
>     ax = jax.exp(a * jax.log(x) - x - jax.lgamma(a))
> 
>     # Power series
>     r = a
>     c = 1.0
>     ans = 1.0
> 
>     while True:
>         r += 1.0
>         c *= x / r
>         ans += c
>         if c / ans <= MACHEP:
>             return ans * ax / a
> 
> 
> def _igamc(a, x):
>     # Compute  x**a * exp(-x) / Gamma(a) ,ax
>     ax = jnp.exp(a * jnp.log(x) - x - jax.lax.lgamma(a))
>     # print(ax)
>     # Continued fraction
>     y = 1.0 - a
>     z = x + y + 1.0
>     c = 0.0
>     pkm2 = 1.0
>     qkm2 = x
>     pkm1 = x + 1.0
>     qkm1 = z * x
>     ans = pkm1 / qkm1
>     while True:
>         c += 1.0
>         y += 1.0
>         z += 2.0
>         yc = y * c
>         pk = pkm1 * z - pkm2 * yc
>         qk = qkm1 * z - qkm2 * yc
>         if qk != 0:
>             r = pk / qk
>             t = abs((ans - r) / r)
>             ans = r
>         else:
>             t = 1.0
>         pkm2 = pkm1
>         pkm1 = pk
>         qkm2 = qkm1
>         qkm1 = qk
>         if abs(pk) > BIG:
>             pkm2 *= BIGINV
>             pkm1 *= BIGINV
>             qkm2 *= BIGINV
>             qkm1 *= BIGINV
>         if t <= MACHEP:
>             return ans * ax
> 
> 
> def chi2(X, y, label_lst):
>     """
>     Calculate the chi-squared statistic and p-value for feature independence testing.
> 
>     Parameters
>     ----------
>     X : array-like, shape (n_samples, n_features)
>         The feature matrix from which to calculate the chi-squared statistic. Each row
>         represents a sample and each column a feature.
> 
>     y : array-like, shape (n_samples,)
>         The class labels of each sample as integers.
> 
>     label_lst : list of str, length = n_classes
>         A list of class labels corresponding to the columns of y.
> 
>     Returns
>     -------
>     chi2_stats : array, shape (n_features,)
>         The chi-squared statistic for each feature, indicating the degree of
>         association between the feature and the class labels.
> 
>     p_value : array, shape (n_features,)
>         The p-value associated with each chi-squared statistic, which can be used
>         to test the null hypothesis that the features are independent of the
>         class labels.
>     """
>     total_samples = X.shape[0]
>     num_feature = X.shape[1]
>     num_class = len(label_lst)
>     # one hot encoding
>     y = jnp.eye(num_class)[y]
>     X = jnp.array(X)
>     y = jnp.array(y)
>     feature_count = jnp.sum(X, axis=0)
>     class_prob = jnp.mean(y, axis=0)
>     # Calculate the observed frequency count
>     observed = jnp.dot(X.T, y)
>     expected = jnp.outer(feature_count, class_prob)
>     # Calculate the chi-squared statistic
>     chi2_stats = (observed - expected) ** 2 / expected
>     # Sum over class dimensions
>     chi2_stats = jnp.nansum(chi2_stats, axis=1)
>     # Degrees of freedom
>     df = num_class - 1
>     # Calculate the p-value for each feature
>     p_value = _sf(chi2_stats, df=df)
>     return chi2_stats, p_value
> 
> 
> if __name__ == '__main__':
>     x, y = load_iris(return_X_y=True)
>     label_lst = np.unique(y)
>     chi2_stats, p_value = chi2(x, y, label_lst)
>     print(chi2_stats, p_value)
>     from sklearn.feature_selection import chi2 as chi2_sklearn
> 
>     sklearn_chi2_stats, sklearn_p_value = chi2_sklearn(x, y)
>     print(sklearn_chi2_stats, sklearn_p_value)
> ```

Hi，明文下的实现会依赖eps (即机器精度)动态决定循环次数，SPU不支持，您可以手动决定一个循环次数上限，然后考察算子精度，只要精度ok即可～
> > 目前的实现如下，按照scipy源码写了一个python版本
> > ```python
> > import jax
> > import jax.numpy as jnp
> > import numpy as np
> > from sklearn.datasets import load_iris
> > 
> > MACHEP = 0.0000001  # the machine roundoff error / tolerance
> > BIG = 4.503599627370496e15
> > BIGINV = 2.22044604925031308085e-16
> > 
> > 
> > def _sf(x, df):
> >     def body_func(x_i):
> >         if x_i < 0:
> >             return 1
> >         if x_i == 0:
> >             return 0
> >         if df <= 0:
> >             raise ValueError("Domain error")
> >         if x_i < 1.0 or x_i < df:
> >             return 1.0 - _igam(0.5 * df, 0.5 * x_i)
> >         return _igamc(df * 0.5, x_i * 0.5)
> > 
> >     return jax.vmap(body_func)(x)
> > 
> > 
> > def _igam(a, x):
> >     # Compute  x**a * exp(-x) / Gamma(a)
> >     ax = jax.exp(a * jax.log(x) - x - jax.lgamma(a))
> > 
> >     # Power series
> >     r = a
> >     c = 1.0
> >     ans = 1.0
> > 
> >     while True:
> >         r += 1.0
> >         c *= x / r
> >         ans += c
> >         if c / ans <= MACHEP:
> >             return ans * ax / a
> > 
> > 
> > def _igamc(a, x):
> >     # Compute  x**a * exp(-x) / Gamma(a) ,ax
> >     ax = jnp.exp(a * jnp.log(x) - x - jax.lax.lgamma(a))
> >     # print(ax)
> >     # Continued fraction
> >     y = 1.0 - a
> >     z = x + y + 1.0
> >     c = 0.0
> >     pkm2 = 1.0
> >     qkm2 = x
> >     pkm1 = x + 1.0
> >     qkm1 = z * x
> >     ans = pkm1 / qkm1
> >     while True:
> >         c += 1.0
> >         y += 1.0
> >         z += 2.0
> >         yc = y * c
> >         pk = pkm1 * z - pkm2 * yc
> >         qk = qkm1 * z - qkm2 * yc
> >         if qk != 0:
> >             r = pk / qk
> >             t = abs((ans - r) / r)
> >             ans = r
> >         else:
> >             t = 1.0
> >         pkm2 = pkm1
> >         pkm1 = pk
> >         qkm2 = qkm1
> >         qkm1 = qk
> >         if abs(pk) > BIG:
> >             pkm2 *= BIGINV
> >             pkm1 *= BIGINV
> >             qkm2 *= BIGINV
> >             qkm1 *= BIGINV
> >         if t <= MACHEP:
> >             return ans * ax
> > 
> > 
> > def chi2(X, y, label_lst):
> >     """
> >     Calculate the chi-squared statistic and p-value for feature independence testing.
> > 
> >     Parameters
> >     ----------
> >     X : array-like, shape (n_samples, n_features)
> >         The feature matrix from which to calculate the chi-squared statistic. Each row
> >         represents a sample and each column a feature.
> > 
> >     y : array-like, shape (n_samples,)
> >         The class labels of each sample as integers.
> > 
> >     label_lst : list of str, length = n_classes
> >         A list of class labels corresponding to the columns of y.
> > 
> >     Returns
> >     -------
> >     chi2_stats : array, shape (n_features,)
> >         The chi-squared statistic for each feature, indicating the degree of
> >         association between the feature and the class labels.
> > 
> >     p_value : array, shape (n_features,)
> >         The p-value associated with each chi-squared statistic, which can be used
> >         to test the null hypothesis that the features are independent of the
> >         class labels.
> >     """
> >     total_samples = X.shape[0]
> >     num_feature = X.shape[1]
> >     num_class = len(label_lst)
> >     # one hot encoding
> >     y = jnp.eye(num_class)[y]
> >     X = jnp.array(X)
> >     y = jnp.array(y)
> >     feature_count = jnp.sum(X, axis=0)
> >     class_prob = jnp.mean(y, axis=0)
> >     # Calculate the observed frequency count
> >     observed = jnp.dot(X.T, y)
> >     expected = jnp.outer(feature_count, class_prob)
> >     # Calculate the chi-squared statistic
> >     chi2_stats = (observed - expected) ** 2 / expected
> >     # Sum over class dimensions
> >     chi2_stats = jnp.nansum(chi2_stats, axis=1)
> >     # Degrees of freedom
> >     df = num_class - 1
> >     # Calculate the p-value for each feature
> >     p_value = _sf(chi2_stats, df=df)
> >     return chi2_stats, p_value
> > 
> > 
> > if __name__ == '__main__':
> >     x, y = load_iris(return_X_y=True)
> >     label_lst = np.unique(y)
> >     chi2_stats, p_value = chi2(x, y, label_lst)
> >     print(chi2_stats, p_value)
> >     from sklearn.feature_selection import chi2 as chi2_sklearn
> > 
> >     sklearn_chi2_stats, sklearn_p_value = chi2_sklearn(x, y)
> >     print(sklearn_chi2_stats, sklearn_p_value)
> > ```
> 
> Hi，明文下的实现会依赖eps (即机器精度)动态决定循环次数，SPU不支持，您可以手动决定一个循环次数上限，然后考察算子精度，只要精度ok即可～

好的，不过可以麻烦您帮我看一下_sf()函数中vmap的使用方法是否正确吗？
我想按照scipy里的思路，根据x和df的大小关系，选择不同的方法（连分数或者级数）来进行计算，但是采用vmap后并不能执行条件判断语句，这方面我不太懂有什么方法去实现。
1. 首先，你这里应该不需要vmap，，x只是一个2维array，处理起来并不复杂。
2. 需要做判断的地方，如选择_igam或者_igamc的时候，可以参考使用`jnp.select`
## 目前的实现
```python
MAX_ITER = 5
def _sf(x, df):
    if df <= 0:
        raise ValueError("Domain error")
    condlist = [x < 0, x == 0, jnp.logical_or(x < 1, x < df)]
    choicelist = [
        jnp.ones_like(x),
        jnp.zeros_like(x),
        1 - _igam(0.5 * df, 0.5 * x),
    ]
    result = jnp.select(condlist, choicelist, default=_igamc(df * 0.5, x * 0.5))
    return result


def _igam(a, x):
    # series
    # Compute  x**a * exp(-x) / Gamma(a)
    ax = jnp.exp(a * jnp.log(x) - x - jax.lax.lgamma(a))
    # Power series
    r = jnp.full_like(x, a)
    c = jnp.ones_like(x)
    ans = jnp.ones_like(x)

    def loop_body(_, val):
        r, c, ans = val
        r += 1.0
        c *= x / r
        ans += c
        return (r, c, ans)

    init_val = (r, c, ans)
    _, _, ans = jax.lax.fori_loop(0, MAX_ITER, loop_body, init_val)
    return ans * ax / a


def _igamc(a, x):
    # Continued fraction
    # Compute ax = x**a * exp(-x) / Gamma(a)
    ax = jnp.exp(a * jnp.log(x) - x - jax.lax.lgamma(a))
    y = jnp.ones_like(x) - a
    z = x + y + 1.0
    c = jnp.zeros_like(x)
    pkm2 = jnp.ones_like(x)
    qkm2 = jnp.full_like(x, a)
    pkm1 = x + 1.0
    qkm1 = z * x
    ans = pkm1 / qkm1

    def loop_body(_, val):
        (c, y, z, pkm1, pkm2, qkm1, qkm2), ans = val
        c += 1.0
        y += 1.0
        z += 2.0
        yc = y * c
        # 连分数的一项
        pk = pkm1 * z - pkm2 * yc
        qk = qkm1 * z - qkm2 * yc
        r = pk / qk
        ans = r
        pkm2 = pkm1
        pkm1 = pk
        qkm2 = qkm1
        qkm1 = qk
        return (c, y, z, pkm1, pkm2, qkm1, qkm2), ans

    params = (c, y, z, pkm1, pkm2, qkm1, qkm2)
    init_val = (params, ans)
    _, ans = jax.lax.fori_loop(0, MAX_ITER, loop_body, init_val)
    return ans * ax
```
## 明文测试结果
<img width="1047" alt="image" src="https://github.com/secretflow/spu/assets/58327725/5c8fc30e-c011-45a7-84d6-def45331baa8">

## SPU测试结果
<img width="1237" alt="image" src="https://github.com/secretflow/spu/assets/58327725/82b89ed0-cc4d-49f2-995d-b9bef69bdf99">

## 问题
目前结果如上，好像还有一些数值溢出问题，请问这部分有没有什么解决思路，已经尝试过降低循环上限，但是后两个值还是算不出来。
> ## 目前的实现
> ```python
> MAX_ITER = 5
> def _sf(x, df):
>     if df <= 0:
>         raise ValueError("Domain error")
>     condlist = [x < 0, x == 0, jnp.logical_or(x < 1, x < df)]
>     choicelist = [
>         jnp.ones_like(x),
>         jnp.zeros_like(x),
>         1 - _igam(0.5 * df, 0.5 * x),
>     ]
>     result = jnp.select(condlist, choicelist, default=_igamc(df * 0.5, x * 0.5))
>     return result
> 
> 
> def _igam(a, x):
>     # series
>     # Compute  x**a * exp(-x) / Gamma(a)
>     ax = jnp.exp(a * jnp.log(x) - x - jax.lax.lgamma(a))
>     # Power series
>     r = jnp.full_like(x, a)
>     c = jnp.ones_like(x)
>     ans = jnp.ones_like(x)
> 
>     def loop_body(_, val):
>         r, c, ans = val
>         r += 1.0
>         c *= x / r
>         ans += c
>         return (r, c, ans)
> 
>     init_val = (r, c, ans)
>     _, _, ans = jax.lax.fori_loop(0, MAX_ITER, loop_body, init_val)
>     return ans * ax / a
> 
> 
> def _igamc(a, x):
>     # Continued fraction
>     # Compute ax = x**a * exp(-x) / Gamma(a)
>     ax = jnp.exp(a * jnp.log(x) - x - jax.lax.lgamma(a))
>     y = jnp.ones_like(x) - a
>     z = x + y + 1.0
>     c = jnp.zeros_like(x)
>     pkm2 = jnp.ones_like(x)
>     qkm2 = jnp.full_like(x, a)
>     pkm1 = x + 1.0
>     qkm1 = z * x
>     ans = pkm1 / qkm1
> 
>     def loop_body(_, val):
>         (c, y, z, pkm1, pkm2, qkm1, qkm2), ans = val
>         c += 1.0
>         y += 1.0
>         z += 2.0
>         yc = y * c
>         # 连分数的一项
>         pk = pkm1 * z - pkm2 * yc
>         qk = qkm1 * z - qkm2 * yc
>         r = pk / qk
>         ans = r
>         pkm2 = pkm1
>         pkm1 = pk
>         qkm2 = qkm1
>         qkm1 = qk
>         return (c, y, z, pkm1, pkm2, qkm1, qkm2), ans
> 
>     params = (c, y, z, pkm1, pkm2, qkm1, qkm2)
>     init_val = (params, ans)
>     _, ans = jax.lax.fori_loop(0, MAX_ITER, loop_body, init_val)
>     return ans * ax
> ```
> 
> ## 明文测试结果
> <img alt="image" width="1047" src="https://private-user-images.githubusercontent.com/58327725/325073740-5c8fc30e-c011-45a7-84d6-def45331baa8.png?jwt=eyJhbGciOiJIUzI1NiIsInR5cCI6IkpXVCJ9.eyJpc3MiOiJnaXRodWIuY29tIiwiYXVkIjoicmF3LmdpdGh1YnVzZXJjb250ZW50LmNvbSIsImtleSI6ImtleTUiLCJleHAiOjE3MTM5MzgxMzEsIm5iZiI6MTcxMzkzNzgzMSwicGF0aCI6Ii81ODMyNzcyNS8zMjUwNzM3NDAtNWM4ZmMzMGUtYzAxMS00NWE3LTg0ZDYtZGVmNDUzMzFiYWE4LnBuZz9YLUFtei1BbGdvcml0aG09QVdTNC1ITUFDLVNIQTI1NiZYLUFtei1DcmVkZW50aWFsPUFLSUFWQ09EWUxTQTUzUFFLNFpBJTJGMjAyNDA0MjQlMkZ1cy1lYXN0LTElMkZzMyUyRmF3czRfcmVxdWVzdCZYLUFtei1EYXRlPTIwMjQwNDI0VDA1NTAzMVomWC1BbXotRXhwaXJlcz0zMDAmWC1BbXotU2lnbmF0dXJlPWJjZTQwOTMzZDgwZjQ4ZGQ4Yzc0NTI0ZjBlNTU3MTM5MTAxZTc0MTAzMzBhNDVmMTg3NDk4MmQ2YTQyOTRhYjcmWC1BbXotU2lnbmVkSGVhZGVycz1ob3N0JmFjdG9yX2lkPTAma2V5X2lkPTAmcmVwb19pZD0wIn0.iH7fOYhRamgmRvTJHxVjFdFecI90dKQfgh043OLRank">
> ## SPU测试结果
> <img alt="image" width="1237" src="https://private-user-images.githubusercontent.com/58327725/325074051-82b89ed0-cc4d-49f2-995d-b9bef69bdf99.png?jwt=eyJhbGciOiJIUzI1NiIsInR5cCI6IkpXVCJ9.eyJpc3MiOiJnaXRodWIuY29tIiwiYXVkIjoicmF3LmdpdGh1YnVzZXJjb250ZW50LmNvbSIsImtleSI6ImtleTUiLCJleHAiOjE3MTM5MzgxMzEsIm5iZiI6MTcxMzkzNzgzMSwicGF0aCI6Ii81ODMyNzcyNS8zMjUwNzQwNTEtODJiODllZDAtY2M0ZC00OWYyLTk5NWQtYjliZWY2OWJkZjk5LnBuZz9YLUFtei1BbGdvcml0aG09QVdTNC1ITUFDLVNIQTI1NiZYLUFtei1DcmVkZW50aWFsPUFLSUFWQ09EWUxTQTUzUFFLNFpBJTJGMjAyNDA0MjQlMkZ1cy1lYXN0LTElMkZzMyUyRmF3czRfcmVxdWVzdCZYLUFtei1EYXRlPTIwMjQwNDI0VDA1NTAzMVomWC1BbXotRXhwaXJlcz0zMDAmWC1BbXotU2lnbmF0dXJlPWVkYTFiZmFlNGZkNjY2NWQxNTEyNjNhNzk1MTEyNWQ4ZWY0NzFhNThhNDE4N2E2YmEwZTYxMDA0ZmFjODAwZDkmWC1BbXotU2lnbmVkSGVhZGVycz1ob3N0JmFjdG9yX2lkPTAma2V5X2lkPTAmcmVwb19pZD0wIn0.0-FAcLQ34kuQhCkJnQi23YhwX-okes2groU_9iIeiRY">
> ## 问题
> 目前结果如上，好像还有一些数值溢出问题，请问这部分有没有什么解决思路，已经尝试过降低循环上限，但是后两个值还是算不出来。

溢出的原因是连分数法中pk和qk的数值增长比较快，可以考虑：
1. 减少MAX_ITER到3 （在这个例子中）
2. 把环改成FM128

建议：
1. MAX_ITER 作为一个超参数，可供用户选择
2. 在chi2的函数文档处说明一下pvalue的计算容易溢出 （e.g. 可以建议用户使用FM128，或者适当减少MAX_ITER）
3. pvalue可以选择性的计算（这虽然和sklearn的行为不同，但我觉得还是有意义的，因为chi统计量的计算非常cheap，但是pvalue计算非常昂贵，而二者的作用几乎是等价的，如果用于feature selection，完全可以只计算前者，不过这个也需要在comment里说明）
4. 有一些简单的优化可以做一下：比如`jnp.exp(a * jnp.log(x) - x - jax.lax.lgamma(a))`，在a不为整数时较高效，但a为整数时，`jnp.power(x, a) * jnp.exp(-x) * jnp.exp(-jax.lax.lgamma(a))` 可以节省一次log算子；（由于a是明文，故这个优化是可以做的）
5. api简化：比如`label_lst`似乎不必要，只需要class的个数即可；另外，关于y的格式也需要做解释，当前实现应该只支持y的取值是0,1,2...这种格式的，这个也需要明确说明出来。
您的建议我已采纳，将max_iter设置为0，环改成fm128。尝试后发现，即使不进入循环，也会发生下溢现象，出现在计算ax语句
```python
ax = jnp.power(x, a) * jnp.exp(-x) * jnp.exp(-jax.lax.lgamma(a))
```
具体是计算`jnp.exp(-x)`时发生了下溢，这里我尝试采用`jnp.exp(-(x-jnp.max(x))`进行修正，但仍旧会出现溢出
下溢是正常的，很小的数直接变成0也合理啊

Fix patch:
diff --git a/sml/feature_selection/BUILD.bazel b/sml/feature_selection/BUILD.bazel
new file mode 100644
index 00000000..dc512dcb
--- /dev/null
+++ b/sml/feature_selection/BUILD.bazel
@@ -0,0 +1,22 @@
+# Copyright 2024 Ant Group Co., Ltd.
+#
+# 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.
+
+load("@rules_python//python:defs.bzl", "py_library")
+
+package(default_visibility = ["//visibility:public"])
+
+py_library(
+    name = "univariate_selection",
+    srcs = ["univariate_selection.py"],
+)
diff --git a/sml/feature_selection/emulations/BUILD.bazel b/sml/feature_selection/emulations/BUILD.bazel
new file mode 100644
index 00000000..1cb01658
--- /dev/null
+++ b/sml/feature_selection/emulations/BUILD.bazel
@@ -0,0 +1,26 @@
+# Copyright 2024 Ant Group Co., Ltd.
+#
+# 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.
+
+load("@rules_python//python:defs.bzl", "py_binary")
+
+package(default_visibility = ["//visibility:public"])
+
+py_binary(
+    name = "chi2_emul",
+    srcs = ["chi2_emul.py"],
+    deps = [
+        "//sml/feature_selection:univariate_selection",
+        "//sml/utils:emulation",
+    ],
+)
diff --git a/sml/feature_selection/emulations/chi2_emul.py b/sml/feature_selection/emulations/chi2_emul.py
new file mode 100644
index 00000000..1c5deb61
--- /dev/null
+++ b/sml/feature_selection/emulations/chi2_emul.py
@@ -0,0 +1,96 @@
+# Copyright 2024 Ant Group Co., Ltd.
+#
+# 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 os
+import sys
+import time
+
+import numpy as np
+from sklearn.datasets import load_iris
+from sklearn.feature_selection import chi2 as chi2_sklearn
+
+# Add the library directory to the path
+sys.path.append(os.path.join(os.path.dirname(__file__), '../../../'))
+import sml.utils.emulation as emulation
+from sml.feature_selection.univariate_selection import chi2
+
+
+def emul_Chi2(mode: emulation.Mode.MULTIPROCESS):
+    print("start chi2 stats emulation")
+
+    def load_data():
+        x, y = load_iris(return_X_y=True)
+        return x, y
+
+    def proc(x, y, num_class, max_iter, compute_p_value):
+        chi2_stats, p_value = chi2(x, y, num_class, max_iter, compute_p_value)
+        return chi2_stats, p_value
+
+    try:
+        # bandwidth and latency only work for docker mode
+        emulator = emulation.Emulator(
+            emulation.CLUSTER_ABY3_3PC, mode, bandwidth=300, latency=20
+        )
+        emulator.up()
+        # load data
+        x, y = load_data()
+        label_lst = np.unique(y)
+        num_class = len(label_lst)
+        max_iter = 1
+        compute_p_value = True
+        # compare with sklearn
+        start_time = time.time()
+        sklearn_chi2_stats, sklearn_p_value = chi2_sklearn(x, y)
+        end_time = time.time()
+        print("=========================")
+        print(f"Running time in SKlearn: {end_time - start_time:.3f}s")
+        print("=========================")
+        # mark these data to be protected in SPU
+        X_spu, y_spu = emulator.seal(x, y)
+        # run
+        start_time = time.time()
+        chi2_stats, p_value = emulator.run(proc, static_argnums=(2, 3, 4))(
+            X_spu, y_spu, num_class, max_iter, compute_p_value
+        )
+        end_time = time.time()
+        print("=========================")
+        print(f"Running time in SPU: {end_time - start_time:.3f}s")
+        print("=========================")
+
+        # print chi2 stats
+        print("Chi2 stats result:")
+        print(chi2_stats)
+        print("Sklearn chi2 stats result:")
+        print(sklearn_chi2_stats)
+        assert np.allclose(
+            sklearn_chi2_stats,
+            sklearn_chi2_stats,
+            rtol=1.0e-5,
+            atol=1.0e-2,
+        )
+        print("P value result:")
+        print(p_value)
+        print("Sklearn p value result:")
+        print(sklearn_p_value)
+        assert np.allclose(
+            p_value, sklearn_p_value, rtol=1.0e-5, atol=1.0e-2, equal_nan=True
+        )
+    except Exception as e:
+        print(e)
+    finally:
+        emulator.down()
+
+
+if __name__ == "__main__":
+    emul_Chi2(emulation.Mode.MULTIPROCESS)
diff --git a/sml/feature_selection/univariate_selection.py b/sml/feature_selection/univariate_selection.py
new file mode 100644
index 00000000..a7c91934
--- /dev/null
+++ b/sml/feature_selection/univariate_selection.py
@@ -0,0 +1,236 @@
+# Copyright 2024 Ant Group Co., Ltd.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#      https://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+import jax
+import jax.numpy as jnp
+
+
+def _sf(x, df, max_iter=3):
+    """
+    Calculates the survival function (1 - CDF) of the chi-squared distribution.
+
+    Parameters
+    ----------
+    x : array-like, shape (n_quantiles,)
+        The quantiles at which to compute the survival function.
+
+    df : int
+        The degrees of freedom of the chi-squared distribution. Must be positive.
+
+    max_iter : int, default=3
+        The maximum number of iterations for the numerical computation.
+
+    Returns
+    -------
+    result : array-like, shape (n_quantiles,)
+        The survival function evaluated at the given quantiles.
+
+    Notes
+    -----
+    The `max_iter` parameter is critical for the numerical stability of the computation. A high value can lead to overflow or underflow errors. Consider using a higher precision data type or reducing `max_iter` to ensure stability.
+
+    References
+    ----------
+    [1] "The Digital Library of Mathematical Functions", dlmf.nist.gov.
+
+    [2] Maddock et. al., "Incomplete Gamma Functions",
+        https://www.boost.org/doc/libs/1_61_0/libs/math/doc/html/math_toolkit/sf_gamma/igamma.html
+    """
+    if df <= 0:
+        raise ValueError("Domain error, df must be positive")
+    condlist = [x < 0, x == 0, x < df]
+    choicelist = [
+        jnp.ones_like(x),
+        jnp.zeros_like(x),
+        1 - _igam(0.5 * df, 0.5 * x, max_iter),
+    ]
+    result = jnp.select(
+        condlist, choicelist, default=_igamc(df * 0.5, x * 0.5, max_iter)
+    )
+    return result
+
+
+def _igam(a, x, max_iter=3):
+    """
+    Computes the regularized lower incomplete gamma function using a power series expansion (DLMF 8.9.2) for a given shape parameter.
+
+    Parameters
+    ----------
+    a : int
+        The shape parameter of the gamma function.
+
+    x : array-like, shape (n_quantiles,)
+        The quantiles at which to compute the incomplete gamma function.
+
+    max_iter : int, default=3
+        The maximum number of iterations for the numerical computation.
+
+    Returns
+    -------
+    ans : array-like, shape (n_quantiles,)
+        The regularized lower incomplete gamma function evaluated at the given quantiles.
+
+    Notes
+    -----
+    The `max_iter` parameter is crucial for the numerical stability of the computation. A high value can lead to overflow or underflow errors. It is recommended to use a higher precision data type or reduce the `max_iter` value to avoid such issues.
+
+    References
+    ----------
+    [1] "The Digital Library of Mathematical Functions", dlmf.nist.gov.
+    """
+    if isinstance(a, int):
+        ax = jnp.power(x, a) * jnp.exp(-x) * jnp.exp(-jax.lax.lgamma(a))
+    else:
+        ax = jnp.exp(a * jnp.log(x) - x - jax.lax.lgamma(a))
+    # Power series
+    r = jnp.full_like(x, a)
+    c = jnp.ones_like(x)
+    ans = jnp.ones_like(x)
+
+    def loop_body(_, val):
+        r, c, ans = val
+        r += 1.0
+        c *= x / r
+        ans += c
+        return (r, c, ans)
+
+    init_val = (r, c, ans)
+    _, _, ans = jax.lax.fori_loop(0, max_iter, loop_body, init_val)
+    return ans * ax / a
+
+
+def _igamc(a, x, max_iter=3):
+    """
+    Computes the complementary regularized lower incomplete gamma function using a continued fraction representation (DLMF 8.11.4) for a given shape parameter.
+
+    Parameters
+    ----------
+    a : int
+        The shape parameter of the gamma function.
+
+    x : array-like, shape (n_quantiles,)
+        The quantiles at which to compute the complementary of incomplete gamma function.
+
+    max_iter : int, default=3
+        The maximum number of iterations for the numerical computation.
+
+    Returns
+    -------
+    ans : array-like, shape (n_quantiles,)
+        The complementary regularized lower incomplete gamma function evaluated at the given quantiles.
+
+    Notes
+    -----
+    The `max_iter` parameter is crucial for the numerical stability of the computation. A high value can lead to overflow or underflow errors. It is recommended to use a higher precision data type or reduce the `max_iter` value to avoid such issues.
+
+    References
+    ----------
+    [1] "The Digital Library of Mathematical Functions", dlmf.nist.gov.
+    """
+    # Compute ax = x**a * exp(-x) / Gamma(a)
+    if isinstance(a, int):
+        ax = jnp.power(x, a) * jnp.exp(-x) * jnp.exp(-jax.lax.lgamma(a))
+    else:
+        ax = jnp.exp(a * jnp.log(x) - x - jax.lax.lgamma(a))
+    y = jnp.ones_like(x) - a
+    z = x + y + 1.0
+    c = jnp.zeros_like(x)
+    pkm2 = jnp.ones_like(x)
+    qkm2 = jnp.full_like(x, a)
+    pkm1 = x + 1.0
+    qkm1 = z * x
+    ans = pkm1 / qkm1
+
+    # Using Continued Fraction
+    def loop_body(_, val):
+        (c, y, z, pkm1, pkm2, qkm1, qkm2), ans = val
+        c += 1.0
+        y += 1.0
+        z += 2.0
+        yc = y * c
+        # pk and qk will grow rapidly, leading to a dependency on the size of the ring.
+        # It may be necessary to use FM128 to get a higher precision.
+        pk = pkm1 * z - pkm2 * yc
+        qk = qkm1 * z - qkm2 * yc
+        r = pk / qk
+        ans = r
+        pkm2 = pkm1
+        pkm1 = pk
+        qkm2 = qkm1
+        qkm1 = qk
+        return (c, y, z, pkm1, pkm2, qkm1, qkm2), ans
+
+    params = (c, y, z, pkm1, pkm2, qkm1, qkm2)
+    init_val = (params, ans)
+    params, ans = jax.lax.fori_loop(0, max_iter, loop_body, init_val)
+    return ans * ax
+
+
+def chi2(X, y, n_classes, max_iter=3, compute_p_value=False):
+    """
+    Performs a chi-squared test for independence between features and class labels.
+
+    Parameters
+    ----------
+    X : array-like, shape (n_samples, n_features)
+        The input feature matrix where each row represents a sample and each column a feature.
+
+    y : array-like, shape (n_samples,)
+        An array of class labels for the samples in `X`. Labels should be integers in the range [0, n_classes-1].
+
+    n_classes : int
+        The number of unique classes in the classification problem.
+
+    max_iter : int, default=3
+        The maximum number of iterations for the numerical computation of the p-value. This parameter is crucial for numerical stability and convergence.
+
+    compute_p_value : bool, default=False
+        Whether to compute the p-value alongside the chi-squared statistic.
+
+    Returns
+    -------
+    chi2_stats : array-like, shape (n_features,)
+        The chi-squared statistics computed for each feature.
+
+    p_value : array-like, shape (n_features,), optional
+        The p-values corresponding to each feature's chi-squared statistic. Returned only if `compute_p_value` is set to True.
+
+    Notes
+    -----
+    - The `max_iter` parameter should be set with consideration of the potential for numerical overflow or underflow. In cases where numerical stability is a concern, consider using a higher precision data type (e.g., Float128 if available) or reducing the `max_iter` value.
+    - The p value computation cost maybe expensive. For applications such as feature selection where only the chi-squared statistic is needed, it may be more efficient to set `compute_p_value` to False.
+    - The input `y` should be provided as a 1-dimensional array with class labels encoded as integers in the specified range.
+    """
+
+    # one hot encoding
+    y = jnp.eye(n_classes)[y]
+    X = jnp.array(X)
+    y = jnp.array(y)
+    feature_count = jnp.sum(X, axis=0)
+    class_prob = jnp.mean(y, axis=0)
+    # Calculate the observed frequency count
+    observed = jnp.dot(X.T, y)
+    expected = jnp.outer(feature_count, class_prob)
+    # Calculate the chi-squared statistic
+    chi2_stats = (observed - expected) ** 2 / expected
+    # Sum over class dimensions
+    chi2_stats = jnp.nansum(chi2_stats, axis=1)
+    # Degrees of freedom
+    df = n_classes - 1
+    p_value = None
+    if compute_p_value:
+        # Calculate the p-value for each feature
+        p_value = _sf(chi2_stats, df=df, max_iter=max_iter)
+        p_value = jnp.where(p_value < 0, 0, p_value)
+    return chi2_stats, p_value


Test patch:
diff --git a/sml/feature_selection/tests/BUILD.bazel b/sml/feature_selection/tests/BUILD.bazel
new file mode 100644
index 00000000..1a776e0b
--- /dev/null
+++ b/sml/feature_selection/tests/BUILD.bazel
@@ -0,0 +1,27 @@
+# Copyright 2024 Ant Group Co., Ltd.
+#
+# 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.
+
+load("@rules_python//python:defs.bzl", "py_test")
+
+package(default_visibility = ["//visibility:public"])
+
+py_test(
+    name = "chi2_test",
+    srcs = ["chi2_test.py"],
+    deps = [
+        "//sml/feature_selection:univariate_selection",
+        "//spu:init",
+        "//spu/utils:simulation",
+    ],
+)
diff --git a/sml/feature_selection/tests/chi2_test.py b/sml/feature_selection/tests/chi2_test.py
new file mode 100644
index 00000000..51435eab
--- /dev/null
+++ b/sml/feature_selection/tests/chi2_test.py
@@ -0,0 +1,69 @@
+# Copyright 2024 Ant Group Co., Ltd.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#      https://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+import os
+import sys
+import unittest
+
+import numpy as np
+import spu.spu_pb2 as spu_pb2
+import spu.utils.simulation as spsim
+from sklearn.datasets import load_iris
+from sklearn.feature_selection import chi2 as chi2_sklearn
+
+sys.path.append(os.path.join(os.path.dirname(__file__), '../../../'))
+from sml.feature_selection.univariate_selection import chi2
+
+
+class UnitTests(unittest.TestCase):
+    def test_chi2(self):
+        sim = spsim.Simulator.simple(
+            3, spu_pb2.ProtocolKind.ABY3, spu_pb2.FieldType.FM128
+        )
+
+        def proc(x, y, num_class, max_iter, compute_p_value):
+            chi2_stats, p_value = chi2(x, y, num_class, max_iter, compute_p_value)
+            return chi2_stats, p_value
+
+        # create dataset
+        x, y = load_iris(return_X_y=True)
+        label_lst = np.unique(y)
+        num_class = len(label_lst)
+        max_iter = 1
+        compute_p_value = True
+        chi2_stats, p_value = spsim.sim_jax(sim, proc, static_argnums=(2, 3, 4))(
+            x, y, num_class, max_iter, compute_p_value
+        )
+        sklearn_chi2_stats, sklearn_p_value = chi2_sklearn(x, y)
+        print("Chi2 stats result:")
+        print(chi2_stats)
+        print("Sklearn chi2 stats result:")
+        print(sklearn_chi2_stats)
+        assert np.allclose(
+            sklearn_chi2_stats,
+            sklearn_chi2_stats,
+            rtol=1.0e-5,
+            atol=1.0e-2,
+        )
+        print("P value result:")
+        print(p_value)
+        print("Sklearn p value result:")
+        print(sklearn_p_value)
+        assert np.allclose(
+            p_value, sklearn_p_value, rtol=1.0e-5, atol=1.0e-2, equal_nan=True
+        )
+
+
+if __name__ == "__main__":
+    unittest.main()