Repository: secretflow/spu (language: python)

Issue:
在 SML 中实现 metric 算子（回归）
> 此 ISSUE 为 [隐语开源共建计划（SecretFlow Open Source Contribution Plan，简称 SF OSCP）第三期](https://github.com/orgs/secretflow/discussions/1007)任务 ISSUE，欢迎社区开发者参与共建～
> 若有感兴趣想要认领的任务，但还未报名，辛苦先完成[报名](https://www.wjx.top/vm/QKio0dq.aspx# )进行哈～

## 任务介绍
- 任务名称：在 SML 中实现 metric 算子(回归)
- 技术方向：SPU/SML
- 任务难度：进阶🌟🌟

## 详细要求
- 为 SML 增加 metric 算子（回归），包括：
    ⅰ. explained_variance_score
    ⅱ. mean_squared_error
    ⅲ. mean_poisson_deviance
    ⅳ. mean_gamma_deviance
    ⅴ. d2_tweedie_score
具体功能可参考[sklearn](https://scikit-learn.org/stable/modules/classes.html#module-sklearn.metrics)
- 正确性：请确保提交的代码内容为可以直接运行的
- 代码规范：Python 代码需要使用 black+isort 进行格式化（流水线包含代码规范检查卡点）; bazel需要使用buildifier格式化
- 一次认领需要实现所有算子
  
若有其他建议实现的算法，也可在本 ISSUE 下回复


## 能力要求
- 熟悉经典的机器学习算法
- 熟悉 JAX 或 NumPy，可以使用 NumPy 实现算法

## 操作说明
- [develop说明](https://github.com/secretflow/spu/blob/main/sml/development.md)
- [参考PR](https://github.com/secretflow/spu/pull/240)

## 认领说明
- 请在认领任务后，在该 issue 下 comment 你的具体设计思路
- 设计思路说明：简单说明计划使用什么算法、什么技术方式进行实现


Discussion / hints:
tarantula-leo Give it to me.
@deadlywing
```
from sklearn import metrics
import jax.numpy as jnp

import spu.spu_pb2 as spu_pb2
import spu.utils.simulation as spsim

def _mean_tweedie_deviance(y_true, y_pred, sample_weight, power):
    p = power
    if p < 0:
        # 'Extreme stable', y any real number, y_pred > 0
        dev = 2 * (
            jnp.power(jnp.maximum(y_true, 0), 2 - p) / ((1 - p) * (2 - p))
            - y_true * jnp.power(y_pred, 1 - p) / (1 - p)
            + jnp.power(y_pred, 2 - p) / (2 - p)
        )
    elif p == 0:
        # Normal distribution, y and y_pred any real number
        dev = (y_true - y_pred) ** 2
    elif p == 1:
        # Poisson distribution
        dev = 2 * (y_true * jnp.log((y_true / y_pred)) - y_true + y_pred)
    elif p == 2:
        # Gamma distribution
        dev = 2 * (jnp.log(y_pred / y_true) + y_true / y_pred - 1)
    else:
        dev = 2 * (
            jnp.power(y_true, 2 - p) / ((1 - p) * (2 - p))
            - y_true * jnp.power(y_pred, 1 - p) / (1 - p)
            + jnp.power(y_pred, 2 - p) / (2 - p)
        )
    return jnp.average(dev, weights=sample_weight)

def d2_tweedie_score(y_true, y_pred, sample_weight=None, power=0):
    numerator = _mean_tweedie_deviance(y_true, y_pred, sample_weight=sample_weight, power=power)
    y_avg = jnp.average(y_true, weights=sample_weight)
    denominator = _mean_tweedie_deviance(y_true, y_avg, sample_weight=sample_weight, power=power)
    return 1 - numerator / denominator

def test_d2_tweedie_score():
    y_true = jnp.array([0.5, 1, 2.5, 7])
    y_pred = jnp.array([1, 1, 5, 3.5])
    sim = spsim.Simulator.simple(
        3, spu_pb2.ProtocolKind.ABY3, spu_pb2.FieldType.FM64
    )
    print(metrics.d2_tweedie_score(y_true, y_pred, sample_weight=jnp.array([1, 1, 1, 1]), power=1))
    print(spsim.sim_jax(sim, d2_tweedie_score, static_argnums=(3, ))(y_true, y_pred, jnp.array([1, 1, 1, 1]), 1))


def explained_variance_score(
    y_true,
    y_pred,
    sample_weight=None,
    multioutput="uniform_average",
):
    y_diff_avg = jnp.average(y_true - y_pred, weights=sample_weight, axis=0)
    numerator = jnp.average(
        (y_true - y_pred - y_diff_avg) ** 2, weights=sample_weight, axis=0
    )

    y_true_avg = jnp.average(y_true, weights=sample_weight, axis=0)
    denominator = jnp.average((y_true - y_true_avg) ** 2, weights=sample_weight, axis=0)
    output_scores = 1 - (numerator / denominator)

    if isinstance(multioutput, str):
        if multioutput == "raw_values":
            # return scores individually
            return output_scores
        elif multioutput == "uniform_average":
            # Passing None as weights to np.average results is uniform mean
            avg_weights = None
        elif multioutput == "variance_weighted":
            avg_weights = denominator
    else:
        avg_weights = multioutput

    return jnp.average(output_scores, weights=avg_weights)

def test_explained_variance_score():
    # y_true = jnp.array([3, -0.5, 2, 7])
    # y_pred = jnp.array([2.5, 0.0, 2, 8])
    y_true = jnp.array([[0.5, 1],[-1, 1],[7, -6]])
    y_pred = jnp.array([[0, 2],[-1, 2],[8, -5]])
    print(metrics.explained_variance_score(y_true, y_pred, multioutput="variance_weighted", force_finite=True))
    sim = spsim.Simulator.simple(
        3, spu_pb2.ProtocolKind.ABY3, spu_pb2.FieldType.FM64
    )
    print(spsim.sim_jax(sim, explained_variance_score, static_argnums=(3, ))(y_true, y_pred, None, "variance_weighted"))

def mean_squared_error(y_true, y_pred, sample_weight=None, multioutput="uniform_average", squared=True):
    output_errors = jnp.average((y_true - y_pred) ** 2, axis=0, weights=sample_weight)

    if not squared:
        output_errors = jnp.sqrt(output_errors)
    if isinstance(multioutput, str):
        if multioutput == "raw_values":
            return output_errors
        elif multioutput == "uniform_average":
            # pass None as weights to np.average: uniform mean
            multioutput = None

    return jnp.average(output_errors, weights=multioutput)

def test_mean_squared_error():
    # y_true = jnp.array([3, -0.5, 2, 7])
    # y_pred = jnp.array([2.5, 0.0, 2, 8])
    y_true = jnp.array([[0.5, 1],[-1, 1],[7, -6]])
    y_pred = jnp.array([[0, 2],[-1, 2],[8, -5]])
    sim = spsim.Simulator.simple(
        3, spu_pb2.ProtocolKind.ABY3, spu_pb2.FieldType.FM64
    )
    print(metrics.mean_squared_error(y_true, y_pred, sample_weight=None, squared=False))
    print(spsim.sim_jax(sim, mean_squared_error, static_argnums=(3, 4))(y_true, y_pred, None, "uniform_average", False))


def mean_poisson_deviance(y_true, y_pred, sample_weight=None):
    return _mean_tweedie_deviance(y_true, y_pred, sample_weight=sample_weight, power=1)

def test_mean_poisson_deviance():
    y_true = jnp.array([2, 0, 1, 4])
    y_pred = jnp.array([0.5, 0.5, 2., 2.])
    print(metrics.mean_poisson_deviance(y_true, y_pred))
    sim = spsim.Simulator.simple(
        3, spu_pb2.ProtocolKind.ABY3, spu_pb2.FieldType.FM64
    )
    print(spsim.sim_jax(sim, mean_poisson_deviance)(y_true, y_pred))


def mean_gamma_deviance(y_true, y_pred, sample_weight=None):
    return _mean_tweedie_deviance(y_true, y_pred, sample_weight=sample_weight, power=2)

def test_mean_gamma_deviance():
    y_true = jnp.array([2, 0.5, 1, 4])
    y_pred = jnp.array([0.5, 0.5, 2., 2.])
    print(metrics.mean_gamma_deviance(y_true, y_pred))
    sim = spsim.Simulator.simple(
        3, spu_pb2.ProtocolKind.ABY3, spu_pb2.FieldType.FM64
    )
    print(spsim.sim_jax(sim, mean_gamma_deviance)(y_true, y_pred))
```
hi，正确性上没什么问题，但是性能上能优化；有几种常见的涉及重复的pattern可以思考一下，如对`d2_tweedie_score`: 
   1). 对某些p，分子和分母的计算是重复的（只和`y_true`有关的项，我不确定编译器是否会自动缓存，最好明确的表示）
   2).类似`jnp.power(y_pred, 1 - p)`和`jnp.power(y_pred, 2 - p)`显然不需要两次power
   3).average的问题，对equal weight的情况(None)，实际上分子和分母只需要调用jnp.sum而不需要jnp.average (可以省一次trunc)

为了性能，可能代码书写上会ugly一些，可以仔细考虑一下代码的具体实现方式哈～～
还有一个，可能比较隐晦，p=2时：`jnp.log(y_pred / y_true)`，直接计算的话，这个pattern需要计算两次除，但由于`y_true`是一致的，实际上只需要计算一次`1/y_true`和两次乘法；
如果要比较不同写法底层实际调用的算子次数和性能，可以考虑打开一些profile的开关，自己实验一下～

simulation可以直接uncomment掉simple函数里的这两行
![image](https://github.com/secretflow/spu/assets/25818621/74eb4f9d-ee47-42cc-a772-5a54fcd32ecc)

> 还有一个，可能比较隐晦，p=2时：`jnp.log(y_pred / y_true)`，直接计算的话，这个pattern需要计算两次除，但由于`y_true`是一致的，实际上只需要计算一次`1/y_true`和两次乘法；

这里怎么理解？
另外average的问题，对equal weight的情况，np库在为None时，调用的是sum函数，应该是不需要在外部显示地对这种情况进行判断？
> > 还有一个，可能比较隐晦，p=2时：`jnp.log(y_pred / y_true)`，直接计算的话，这个pattern需要计算两次除，但由于`y_true`是一致的，实际上只需要计算一次`1/y_true`和两次乘法；
> 
> 这里怎么理解？

分子分母需要两次计算`jnp.log(y_pred / y_true)`，其中y_true是一样的，所以可以只计算一次`1/y_true`（除法相对更costly）
> 另外average的问题，对equal weight的情况，np库在为None时，调用的是sum函数，应该是不需要在外部显示地对这种情况进行判断？

我看的jax的实现，似乎是使用mean函数，应该还是有的，，具体你可以实验一下看看具体的MPC层算子调用情况
![image](https://github.com/secretflow/spu/assets/25818621/16d5c548-4e08-4503-8ed0-d5a169b8c5b8)


Fix patch:
diff --git a/sml/metrics/regression/BUILD.bazel b/sml/metrics/regression/BUILD.bazel
new file mode 100644
index 00000000..88961748
--- /dev/null
+++ b/sml/metrics/regression/BUILD.bazel
@@ -0,0 +1,41 @@
+# Copyright 2023 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", "py_test")
+
+package(default_visibility = ["//visibility:public"])
+
+py_library(
+    name = "regression",
+    srcs = ["regression.py"],
+)
+
+py_test(
+    name = "regression_test",
+    srcs = ["regression_test.py"],
+    deps = [
+        ":regression",
+        "//spu:init",
+        "//spu/utils:simulation",
+    ],
+)
+
+py_binary(
+    name = "regression_emul",
+    srcs = ["regression_emul.py"],
+    deps = [
+        ":regression",
+        "//sml/utils:emulation",
+    ],
+)
diff --git a/sml/metrics/regression/regression.py b/sml/metrics/regression/regression.py
new file mode 100644
index 00000000..55a3f0e3
--- /dev/null
+++ b/sml/metrics/regression/regression.py
@@ -0,0 +1,114 @@
+# Copyright 2023 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 specifi
+
+import jax.numpy as jnp
+
+
+def _mean_tweedie_deviance(y_true, y_pred, sample_weight, power, d2_score=False):
+    p = power
+    if p < 0:
+        # 'Extreme stable', y any real number, y_pred > 0
+        temp = 1 - p
+        temp_power = jnp.power(y_pred, temp)
+        dev = 2 * (
+            jnp.power(jnp.maximum(y_true, 0), 1 + temp) / (temp * (temp + 1))
+            - temp_power * (y_true / temp - y_pred / (temp + 1))
+        )
+    elif p == 0:
+        # Normal distribution, y and y_pred any real number
+        dev = (y_true - y_pred) ** 2
+    elif p == 1:
+        # Poisson distribution
+        dev = 2 * (y_true * jnp.log((y_true / y_pred)) - y_true + y_pred)
+    elif p == 2:
+        # Gamma distribution
+        dev = 2 * (jnp.log(y_pred / y_true) + y_true / y_pred - 1)
+    else:
+        temp = 1 - p
+        temp_power = jnp.power(y_pred, temp)
+        dev = 2 * (
+            jnp.power(y_true, 1 + temp) / (temp * (temp + 1))
+            - temp_power * (y_true / temp - y_pred / (temp + 1))
+        )
+
+    if d2_score and (sample_weight is None):
+        # When weight is none and the d2 fraction is calculated, the numerator and denominator are divided by shape, and only sum is used
+        return jnp.sum(dev)
+    else:
+        return jnp.average(dev, weights=sample_weight)
+
+
+def d2_tweedie_score(y_true, y_pred, sample_weight=None, power=0):
+    numerator = _mean_tweedie_deviance(
+        y_true, y_pred, sample_weight=sample_weight, power=power, d2_score=True
+    )
+    y_avg = jnp.average(y_true, weights=sample_weight)
+    denominator = _mean_tweedie_deviance(
+        y_true, y_avg, sample_weight=sample_weight, power=power, d2_score=True
+    )
+    return 1 - numerator / denominator
+
+
+def explained_variance_score(
+    y_true,
+    y_pred,
+    sample_weight=None,
+    multioutput="uniform_average",
+):
+    y_diff_avg = jnp.average(y_true - y_pred, weights=sample_weight, axis=0)
+    numerator = jnp.average(
+        (y_true - y_pred - y_diff_avg) ** 2, weights=sample_weight, axis=0
+    )
+
+    y_true_avg = jnp.average(y_true, weights=sample_weight, axis=0)
+    denominator = jnp.average((y_true - y_true_avg) ** 2, weights=sample_weight, axis=0)
+    output_scores = 1 - (numerator / denominator)
+
+    if isinstance(multioutput, str):
+        if multioutput == "raw_values":
+            # return scores individually
+            return output_scores
+        elif multioutput == "uniform_average":
+            # Passing None as weights to np.average results is uniform mean
+            avg_weights = None
+        elif multioutput == "variance_weighted":
+            avg_weights = denominator
+    else:
+        avg_weights = multioutput
+
+    return jnp.average(output_scores, weights=avg_weights)
+
+
+def mean_squared_error(
+    y_true, y_pred, sample_weight=None, multioutput="uniform_average", squared=True
+):
+    output_errors = jnp.average((y_true - y_pred) ** 2, axis=0, weights=sample_weight)
+
+    if not squared:
+        output_errors = jnp.sqrt(output_errors)
+    if isinstance(multioutput, str):
+        if multioutput == "raw_values":
+            return output_errors
+        elif multioutput == "uniform_average":
+            # pass None as weights to np.average: uniform mean
+            multioutput = None
+
+    return jnp.average(output_errors, weights=multioutput)
+
+
+def mean_poisson_deviance(y_true, y_pred, sample_weight=None):
+    return _mean_tweedie_deviance(y_true, y_pred, sample_weight=sample_weight, power=1)
+
+
+def mean_gamma_deviance(y_true, y_pred, sample_weight=None):
+    return _mean_tweedie_deviance(y_true, y_pred, sample_weight=sample_weight, power=2)
diff --git a/sml/metrics/regression/regression_emul.py b/sml/metrics/regression/regression_emul.py
new file mode 100644
index 00000000..4448c5b9
--- /dev/null
+++ b/sml/metrics/regression/regression_emul.py
@@ -0,0 +1,129 @@
+# Copyright 2023 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 specifi
+
+import os
+import sys
+
+import jax.numpy as jnp
+import numpy as np
+from sklearn import metrics
+
+# add ops dir to the path
+sys.path.append(os.path.join(os.path.dirname(__file__), '../../../'))
+
+
+import sml.utils.emulation as emulation
+from sml.metrics.regression.regression import (
+    explained_variance_score,
+    mean_squared_error,
+    mean_poisson_deviance,
+    mean_gamma_deviance,
+    d2_tweedie_score,
+)
+
+
+def emul_d2_tweedie_score():
+    power_list = [-1, 0, 1, 2, 3]
+    weight_list = [None, jnp.array([0.5, 0.5, 0.5, 0.5]), jnp.array([0.5, 1, 2, 0.5])]
+
+    # Test d2_tweedie_score
+    y_true = jnp.array([0.5, 1, 2.5, 7])
+    y_pred = jnp.array([1, 1, 5, 3.5])
+    for p in power_list:
+        for weight in weight_list:
+            sk_result = metrics.d2_tweedie_score(
+                y_true, y_pred, sample_weight=weight, power=p
+            )
+            spu_result = emulator.run(d2_tweedie_score, static_argnums=(3,))(
+                y_true, y_pred, weight, p
+            )
+            np.testing.assert_allclose(sk_result, spu_result, rtol=0, atol=1e-4)
+
+
+def emul_explained_variance_score():
+    weight_list = [None, jnp.array([0.5, 0.5, 0.5, 0.5]), jnp.array([0.5, 1, 2, 0.5])]
+
+    # Test explained_variance_score
+    y_true = jnp.array([3, -0.5, 2, 7])
+    y_pred = jnp.array([2.5, 0.0, 2, 8])
+    for weight in weight_list:
+        sk_result = metrics.explained_variance_score(
+            y_true,
+            y_pred,
+            sample_weight=weight,
+            multioutput="variance_weighted",
+            force_finite=True,
+        )
+        spu_result = emulator.run(explained_variance_score, static_argnums=(3,))(
+            y_true, y_pred, weight, "variance_weighted"
+        )
+        np.testing.assert_allclose(sk_result, spu_result, rtol=0, atol=1e-4)
+
+
+def emul_mean_squared_error():
+    weight_list = [None, jnp.array([0.5, 0.5, 0.5, 0.5]), jnp.array([0.5, 1, 2, 0.5])]
+
+    # Test mean_squared_error
+    y_true = jnp.array([3, -0.5, 2, 7])
+    y_pred = jnp.array([2.5, 0.0, 2, 8])
+    for weight in weight_list:
+        sk_result = metrics.mean_squared_error(
+            y_true, y_pred, sample_weight=weight, squared=False
+        )
+        spu_result = emulator.run(mean_squared_error, static_argnums=(3, 4))(
+            y_true, y_pred, weight, "uniform_average", False
+        )
+        np.testing.assert_allclose(sk_result, spu_result, rtol=0, atol=1e-4)
+
+
+def emul_mean_poisson_deviance():
+    weight_list = [None, jnp.array([0.5, 0.5, 0.5, 0.5]), jnp.array([0.5, 1, 2, 0.5])]
+
+    # Test mean_poisson_deviance
+    y_true = jnp.array([2, 0, 1, 4])
+    y_pred = jnp.array([0.5, 0.5, 2.0, 2.0])
+    for weight in weight_list:
+        sk_result = metrics.mean_poisson_deviance(y_true, y_pred, sample_weight=weight)
+        spu_result = emulator.run(mean_poisson_deviance)(y_true, y_pred, weight)
+        np.testing.assert_allclose(sk_result, spu_result, rtol=0, atol=1e-4)
+
+
+def emul_mean_gamma_deviance():
+    weight_list = [None, jnp.array([0.5, 0.5, 0.5, 0.5]), jnp.array([0.5, 1, 2, 0.5])]
+
+    # Test mean_gamma_deviance
+    y_true = jnp.array([2, 0.5, 1, 4])
+    y_pred = jnp.array([0.5, 0.5, 2.0, 2.0])
+    for weight in weight_list:
+        sk_result = metrics.mean_gamma_deviance(y_true, y_pred, sample_weight=weight)
+        spu_result = emulator.run(mean_gamma_deviance)(y_true, y_pred, weight)
+        np.testing.assert_allclose(sk_result, spu_result, rtol=0, atol=1e-4)
+
+
+if __name__ == "__main__":
+    try:
+        # bandwidth and latency only work for docker mode
+        emulator = emulation.Emulator(
+            emulation.CLUSTER_ABY3_3PC,
+            emulation.Mode.MULTIPROCESS,
+            bandwidth=300,
+            latency=20,
+        )
+        emulator.up()
+        emul_d2_tweedie_score()
+        emul_explained_variance_score()
+        emul_mean_squared_error()
+        emul_mean_poisson_deviance()
+        emul_mean_gamma_deviance()
+    finally:
+        emulator.down()


Test patch:
diff --git a/sml/metrics/regression/regression_test.py b/sml/metrics/regression/regression_test.py
new file mode 100644
index 00000000..fc868c9e
--- /dev/null
+++ b/sml/metrics/regression/regression_test.py
@@ -0,0 +1,159 @@
+# Copyright 2023 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 specifi
+
+import os
+import sys
+import unittest
+
+import jax.numpy as jnp
+import numpy as np
+from sklearn import metrics
+
+import spu.spu_pb2 as spu_pb2
+import spu.utils.simulation as spsim
+
+# add ops dir to the path
+sys.path.append(os.path.join(os.path.dirname(__file__), '../../../'))
+
+from sml.metrics.regression.regression import (
+    explained_variance_score,
+    mean_squared_error,
+    mean_poisson_deviance,
+    mean_gamma_deviance,
+    d2_tweedie_score,
+)
+
+
+class UnitTests(unittest.TestCase):
+    def test_d2_tweedie_score(self):
+        sim = spsim.Simulator.simple(
+            3, spu_pb2.ProtocolKind.ABY3, spu_pb2.FieldType.FM128
+        )
+
+        power_list = [-1, 0, 1, 2, 3]
+        weight_list = [
+            None,
+            jnp.array([0.5, 0.5, 0.5, 0.5]),
+            jnp.array([0.5, 1, 2, 0.5]),
+        ]
+
+        # Test d2_tweedie_score
+        y_true = jnp.array([0.5, 1, 2.5, 7])
+        y_pred = jnp.array([1, 1, 5, 3.5])
+        for p in power_list:
+            for weight in weight_list:
+                sk_result = metrics.d2_tweedie_score(
+                    y_true, y_pred, sample_weight=weight, power=p
+                )
+                spu_result = spsim.sim_jax(sim, d2_tweedie_score, static_argnums=(3,))(
+                    y_true, y_pred, weight, p
+                )
+                np.testing.assert_allclose(sk_result, spu_result, rtol=0, atol=1e-4)
+
+    def test_explained_variance_score(self):
+        sim = spsim.Simulator.simple(
+            3, spu_pb2.ProtocolKind.ABY3, spu_pb2.FieldType.FM128
+        )
+
+        weight_list = [
+            None,
+            jnp.array([0.5, 0.5, 0.5, 0.5]),
+            jnp.array([0.5, 1, 2, 0.5]),
+        ]
+
+        # Test explained_variance_score
+        y_true = jnp.array([3, -0.5, 2, 7])
+        y_pred = jnp.array([2.5, 0.0, 2, 8])
+        for weight in weight_list:
+            sk_result = metrics.explained_variance_score(
+                y_true,
+                y_pred,
+                sample_weight=weight,
+                multioutput="variance_weighted",
+                force_finite=True,
+            )
+            spu_result = spsim.sim_jax(
+                sim, explained_variance_score, static_argnums=(3,)
+            )(y_true, y_pred, weight, "variance_weighted")
+            np.testing.assert_allclose(sk_result, spu_result, rtol=0, atol=1e-4)
+
+    def test_mean_squared_error(self):
+        sim = spsim.Simulator.simple(
+            3, spu_pb2.ProtocolKind.ABY3, spu_pb2.FieldType.FM128
+        )
+
+        weight_list = [
+            None,
+            jnp.array([0.5, 0.5, 0.5, 0.5]),
+            jnp.array([0.5, 1, 2, 0.5]),
+        ]
+
+        # Test mean_squared_error
+        y_true = jnp.array([3, -0.5, 2, 7])
+        y_pred = jnp.array([2.5, 0.0, 2, 8])
+        for weight in weight_list:
+            sk_result = metrics.mean_squared_error(
+                y_true, y_pred, sample_weight=weight, squared=False
+            )
+            spu_result = spsim.sim_jax(sim, mean_squared_error, static_argnums=(3, 4))(
+                y_true, y_pred, weight, "uniform_average", False
+            )
+            np.testing.assert_allclose(sk_result, spu_result, rtol=0, atol=1e-4)
+
+    def test_mean_poisson_deviance(self):
+        sim = spsim.Simulator.simple(
+            3, spu_pb2.ProtocolKind.ABY3, spu_pb2.FieldType.FM128
+        )
+
+        weight_list = [
+            None,
+            jnp.array([0.5, 0.5, 0.5, 0.5]),
+            jnp.array([0.5, 1, 2, 0.5]),
+        ]
+
+        # Test mean_poisson_deviance
+        y_true = jnp.array([2, 0, 1, 4])
+        y_pred = jnp.array([0.5, 0.5, 2.0, 2.0])
+        for weight in weight_list:
+            sk_result = metrics.mean_poisson_deviance(
+                y_true, y_pred, sample_weight=weight
+            )
+            spu_result = spsim.sim_jax(sim, mean_poisson_deviance)(
+                y_true, y_pred, weight
+            )
+            np.testing.assert_allclose(sk_result, spu_result, rtol=0, atol=1e-4)
+
+    def test_mean_gamma_deviance(self):
+        sim = spsim.Simulator.simple(
+            3, spu_pb2.ProtocolKind.ABY3, spu_pb2.FieldType.FM128
+        )
+
+        weight_list = [
+            None,
+            jnp.array([0.5, 0.5, 0.5, 0.5]),
+            jnp.array([0.5, 1, 2, 0.5]),
+        ]
+
+        # Test mean_gamma_deviance
+        y_true = jnp.array([2, 0.5, 1, 4])
+        y_pred = jnp.array([0.5, 0.5, 2.0, 2.0])
+        for weight in weight_list:
+            sk_result = metrics.mean_gamma_deviance(
+                y_true, y_pred, sample_weight=weight
+            )
+            spu_result = spsim.sim_jax(sim, mean_gamma_deviance)(y_true, y_pred, weight)
+            np.testing.assert_allclose(sk_result, spu_result, rtol=0, atol=1e-4)
+
+
+if __name__ == "__main__":
+    unittest.main()