Repository: secretflow/spu (language: python)

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


## 任务介绍

- 任务名称：使用 SPU 实现高斯过程分类算法
- 技术方向：SPU/SML
- 任务难度：进阶🌟🌟

## 详细要求
- 安全性（尽量少 reveal）
- 功能性：实现算法的基本功能，包括：
  - 支持 fit，predict
  - 支持起码一种 fit 算法
  - 起码支持一种 kernel（如 RBF）
- 收敛性：包含 simulator 跑出的实验数据并且证明收敛性/准确性（最好有与明文 sklearn 结果的对比，可详见[参考 PR](https://github.com/secretflow/spu/pull/240)）
- 正确性：请确保提交的代码内容为可以直接运行的
- 代码规范：Python 代码需要使用 black+isort 进行格式化（流水线包含代码规范检查卡点）
- 提交说明：关联该 isuue 并提交代码至 https://github.com/secretflow/spu/tree/main/sml
- 特殊说明：若某个特性有特殊的限制，如需要 FM128，需要更多 fxp 等需要在注释文档中明确说明

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

## 操作说明
- 操作指引：https://www.secretflow.org.cn/docs/spu/latest/en-US/getting_started/tutorials/develop_your_first_mpc_application
- 可参考明文实现：
  - [sklearn](https://github.com/scikit-learn/scikit-learn/blob/7f9bad99d/sklearn/gaussian_process/_gpc.py#L487)
  - [ref book](http://gaussianprocess.org/gpml/chapters/)
-  [范例](https://github.com/secretflow/spu/blob/main/sml/linear_model/simple_sgd.py)
-  [参考 PR](https://github.com/secretflow/spu/pull/240)


## 认领说明
本任务可有多种实现方式，故支持一个任务有多位开发者进行认领，请在认领任务后，在该 issue 下 comment 你的具体设计思路。

**设计思路说明**：简单说明计划使用什么算法 or 什么优化器实现任务需求即可

> 当同一任务有多位开发者提交设计思路时：
> -  若多位开发者的设计思路均类似，则将按照 comment 的时间 assign 给第一位 comment 的开发者；
> -  若多位开发者的设计思路均不同，则该 issue 将会拆分为以不同设计思路进行实现的子 issue，并 assign 给对应开发者。

## 开发须知
以下部分代码请必须增加代码注释，对对应代码模块进行说明，包括：
-  __init__函数的超参数含义
-  fit 的具体算法实现说明



Discussion / hints:
salazar1117 Give it to me.
初始化：指定核函数kernal（及所需参数）
优化器：L-BFGS-B
核函数：RBF核函数（拥有超参数sigma）
fit方法：输入为训练集特征矩阵X和标签y（y至少包含两个不同的值），先用核函数得到协方差矩阵K，在主对角线元素上添加噪声方差，之后对其进行Cholesky分解得到L矩阵。将y标准化为均值为0、方差为1的形式y‘，使用标准化后的y’和L矩阵求解得到L*L^T*alpha=y’中的alpha。过程中得到模型的alpha、L、y的方差与均值，存放在类成员变量中。
predict：对测试样本x，得到其与训练矩阵X的核函数，其转置乘以(y-训练集上预测的概率)得到测试集的预测值并输出预测。
LGTM
> salazar1117 Give it to me. 初始化：指定核函数kernal（及所需参数） 优化器：L-BFGS-B 核函数：RBF核函数（拥有超参数sigma） fit方法：输入为训练集特征矩阵X和标签y（y至少包含两个不同的值），先用核函数得到协方差矩阵K，在主对角线元素上添加噪声方差，之后对其进行Cholesky分解得到L矩阵。将y标准化为均值为0、方差为1的形式y‘，使用标准化后的y’和L矩阵求解得到L_L^T_alpha=y’中的alpha。过程中得到模型的alpha、L、y的方差与均值，存放在类成员变量中。 predict：对测试样本x，得到其与训练矩阵X的核函数，其转置乘以(y-训练集上预测的概率)得到测试集的预测值并输出预测。

LGTM 感谢您的设计方案。
PS：因目前的设计方案已基本涵盖可实现的算法与分布形式，故本 issue 不再拆分，直接 assign 给这位开发者。

Fix patch:
diff --git a/sml/gaussian_process/BUILD.bazel b/sml/gaussian_process/BUILD.bazel
index 7832e732..9b700075 100644
--- a/sml/gaussian_process/BUILD.bazel
+++ b/sml/gaussian_process/BUILD.bazel
@@ -11,3 +11,22 @@
 # 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 = "_gpc",
+    srcs = ["_gpc.py"],
+)
+
+py_library(
+    name = "kernels",
+    srcs = ["kernels.py"],
+)
+
+py_library(
+    name = "ovo_ovr",
+    srcs = ["ovo_ovr.py"],
+)
diff --git a/sml/gaussian_process/_gpc.py b/sml/gaussian_process/_gpc.py
new file mode 100644
index 00000000..66e4ff00
--- /dev/null
+++ b/sml/gaussian_process/_gpc.py
@@ -0,0 +1,318 @@
+import os
+import sys
+
+import jax
+import jax.numpy as jnp
+from jax import grad
+from jax.lax.linalg import cholesky
+from jax.scipy.linalg import cho_solve, solve
+from jax.scipy.special import erf, expit
+
+sys.path.append(os.path.join(os.path.dirname(__file__), "./"))
+from kernels import RBF
+from ovo_ovr import OneVsRestClassifier
+
+LAMBDAS = jnp.array([0.41, 0.4, 0.37, 0.44, 0.39])[:, jnp.newaxis]
+COEFS = jnp.array(
+    [-1854.8214151, 3516.89893646, 221.29346712, 128.12323805, -2010.49422654]
+)[:, jnp.newaxis]
+
+
+class _BinaryGaussianProcessClassifierLaplace:
+    def __init__(
+        self,
+        kernel=None,
+        *,
+        poss="sigmoid",
+        max_iter_predict=100,
+    ):
+        self.kernel = kernel
+        self.max_iter_predict = max_iter_predict
+        self.poss = poss
+
+    def fit(self, X, y):
+        self._check_kernal()
+
+        self.X_train_ = jnp.asarray(X)
+
+        if self.poss == "sigmoid":
+            self.approx_func = expit
+        else:
+            raise ValueError(
+                f"Unsupported prior-likelihood function {self.poss}."
+                "Please try the default dunction sigmoid."
+            )
+
+        self.y_train = y
+
+        K = self.kernel_(self.X_train_)
+        self.f_ = self._posterior_mode(K)
+        return self.f_
+
+    # def log_and_grad(self, f, y_train):
+    #     _tmp = lambda f, y_train: jnp.sum(self.approx_func(y_train*f))
+    #     return grad(_tmp)(f, y_train)/self.approx_func(y_train*f)
+
+    # def log_and_2grads_and_negtive(self, f, y_train):
+    #     _tmp = lambda f, y: jnp.sum(self.log_and_grad(f, y))
+    #     return -grad(_tmp)(f, y_train)
+
+    # def log_and_3grads(self, f, y_train):
+    #     _tmp = lambda f, y_train: jnp.sum(-self.log_and_2grads_and_negtive(f, y_train))
+    #     return grad(_tmp)(f, y_train)
+
+    def predict(self, Xll):
+        X = jnp.asarray(Xll)
+        K_star = self.kernel_(self.X_train_, X)
+        # f_star = K_star.T.dot(self.log_and_grad(self.f_, self.y_train))
+        f_star = K_star.T.dot(self.y_train - self.approx_func(self.f_))
+
+        return jnp.where(f_star > 0, 1, 0)
+
+    def predict_proba(self, Xll):
+        X = jnp.asarray(Xll)
+        K = self.kernel_(self.X_train_)
+        # Based on Algorithm 3.2 of GPML
+
+        # W = self.log_and_2grads_and_negtive(self.f_, self.y_train)
+        pi = self.approx_func(self.f_)
+        W = pi * (1 - pi)
+
+        W_sqr = jnp.sqrt(W)
+        W_sqr_K = W_sqr[:, jnp.newaxis] * K
+        B = jnp.eye(W.shape[0]) + W_sqr_K * W_sqr
+        L = cholesky(B)
+
+        K_star = self.kernel_(self.X_train_, X)
+        # f_star = K_star.T.dot(self.log_and_grad(self.f_, self.y_train))
+        f_star = K_star.T.dot(self.y_train - pi)
+        v = solve(L, W_sqr[:, jnp.newaxis] * K_star)
+        var_f_star = jnp.diag(self.kernel_(X)) - jnp.einsum("ij,ij->j", v, v)
+
+        alpha = 1 / (2 * var_f_star)
+        gamma = LAMBDAS * f_star
+        integrals = (
+            jnp.sqrt(jnp.pi / alpha)
+            * erf(gamma * jnp.sqrt(alpha / (alpha + LAMBDAS**2)))
+            / (2 * jnp.sqrt(var_f_star * 2 * jnp.pi))
+        )
+        pi_star = (COEFS * integrals).sum(axis=0) + 0.5 * COEFS.sum()
+
+        return jnp.vstack((1 - pi_star, pi_star)).T
+
+    def _posterior_mode(self, K):
+        # Based on Algorithm 3.1 of GPML
+        f = jnp.zeros_like(
+                self.y_train, dtype=jnp.float32
+            )  # a warning is triggered if float64 is used
+
+        for _ in range(self.max_iter_predict):
+            # W = self.log_and_2grads_and_negtive(f, self.y_train)
+            pi = self.approx_func(f)
+            W = pi * (1 - pi)
+            W_sqr = jnp.sqrt(W)
+            W_sqr_K = W_sqr[:, jnp.newaxis] * K
+
+            B = jnp.eye(W.shape[0]) + W_sqr_K * W_sqr
+            L = cholesky(B)
+            # b = W * f + self.log_and_grad(f, self.y_train)
+            b = W * f + (self.y_train - pi)
+            a = b - jnp.dot(
+                W_sqr[:, jnp.newaxis] * cho_solve((L, True), jnp.eye(W.shape[0])),
+                W_sqr_K.dot(b),
+            )
+            f = K.dot(a)
+
+        self.f_cached = f  # for warm-start
+        return f
+
+    def _check_kernal(self):
+        if self.kernel is None:  # Use an RBF kernel as default
+            self.kernel_ = RBF()
+        else:
+            self.kernel_ = self.kernel
+
+
+class GaussianProcessClassifier:
+    """Gaussian process classification (GPC) based on Laplace approximation.
+
+    The implementation is based on Algorithm 3.1, 3.2, and 5.1 from [RW2006]_.
+
+    Internally, the Laplace approximation is used for approximating the
+    non-Gaussian posterior by a Gaussian.
+
+    Currently, the implementation is restricted to using the logistic link
+    function. For multi-class classification, several binary one-versus rest
+    classifiers are fitted. Note that this class thus does not implement
+    a true multi-class Laplace approximation.
+
+    Read more in the :ref:`User Guide <gaussian_process>`.
+
+    Parameters
+    ----------
+    kernel : kernel instance, default=None
+        The kernel specifying the covariance function of the GP. If None is
+        passed, the kernel "1.0 * RBF(1.0)" is used as default. Note that
+        the kernel's hyperparameters are optimized during fitting. Also kernel
+        cannot be a `CompoundKernel`.
+
+    max_iter_predict : int, default=100
+        The maximum number of iterations in Newton's method for approximating
+        the posterior during predict. Smaller values will reduce computation
+        time at the cost of worse results.
+
+    multi_class : 'one_vs_rest', default='one_vs_rest'
+        Specifies how multi-class classification problems are handled.
+        One binary Gaussian process classifier is fitted for each class, which
+        is trained to separate this class from the rest.
+
+    poss : "sigmoid", allable or None, default="sigmoid", the predefined 
+        likelihood function which computes the possibility of the predict output 
+        w.r.t. f value.
+
+    Attributes
+    ----------
+    base_estimator_ : ``Estimator`` instance
+        The estimator instance that defines the likelihood function
+        using the observed data.
+
+    kernel_ : kernel instance
+        The kernel used for prediction. In case of binary classification,
+        the structure of the kernel is the same as the one passed as parameter
+        but with optimized hyperparameters. In case of multi-class
+        classification, a CompoundKernel is returned which consists of the
+        different kernels used in the one-versus-rest classifiers.
+
+    n_classes_ : int
+        The number of classes in the training data
+
+    References
+    ----------
+    .. [RW2006] `Carl E. Rasmussen and Christopher K.I. Williams,
+       "Gaussian Processes for Machine Learning",
+       MIT Press 2006 <https://www.gaussianprocess.org/gpml/chapters/RW.pdf>`_
+
+    Examples
+    --------
+    >>> from sklearn.datasets import load_iris
+    >>> from sklearn.gaussian_process import GaussianProcessClassifier
+    >>> from sklearn.gaussian_process.kernels import RBF
+    >>> X, y = load_iris(return_X_y=True)
+    >>> kernel = 1.0 * RBF(1.0)
+    >>> gpc = GaussianProcessClassifier(kernel=kernel,
+    ...         random_state=0).fit(X, y, 3)
+    >>> gpc.predict_proba(X[:2,:])
+    array([[0.83548752, 0.03228706, 0.13222543],
+           [0.79064206, 0.06525643, 0.14410151]])
+    """
+
+    def __init__(
+        self,
+        kernel=None,
+        *,
+        poss="sigmoid",
+        max_iter_predict=100,
+        multi_class="one_vs_rest",
+    ):
+        self.kernel = kernel
+        self.max_iter_predict = max_iter_predict
+        self.multi_class = multi_class
+        self.poss = poss
+
+    def fit(self, X, y, n_classes_):
+        """Fit Gaussian process classification model.
+
+        Parameters
+        ----------
+        X : jax numpy array (n_samples, n_features) of object
+            Feature vectors of training data.
+
+        y : jax numpy array (n_samples,) Target values, 
+        must be preprocessed to 0, 1, 2, ...
+
+        n_classes_ : The number of classes in the training data
+
+        Returns
+        -------
+        self : object
+            Returns an instance of self.
+        """
+        self.n_classes_ = n_classes_
+        self.y_train = jnp.array(y)
+
+        if self.n_classes_ == 1:
+            raise ValueError(
+                "GaussianProcessClassifier requires 2 or more "
+                "distinct classes; got 1 class (only class %s "
+                "is present)" % self.n_classes_[0]
+            )
+
+        self.base_estimator_ = _BinaryGaussianProcessClassifierLaplace(
+            kernel=self.kernel,
+            max_iter_predict=self.max_iter_predict,
+            poss=self.poss,
+        )
+
+        if self.n_classes_ > 2:
+            if self.multi_class == "one_vs_rest":
+                self.base_estimator_ = OneVsRestClassifier(
+                    self.base_estimator_, self.n_classes_
+                )
+            elif self.multi_class == "one_vs_one":
+                raise ValueError("one_vs_one classifier is not supported")
+            else:
+                raise ValueError("Unknown multi-class mode %s" % self.multi_class)
+
+        self.X = jnp.array(X)
+        self.base_estimator_.fit(self.X, self.y_train)
+
+    def predict(self, X):
+        """Perform classification on an array of test vectors X.
+
+        Parameters
+        ----------
+        X : jax numpy array (n_samples, n_features) of object
+            Query points where the GP is evaluated for classification.
+
+        Returns
+        -------
+        C : jax numpy array (n_samples,)
+            Predicted target values for X.
+        """
+        self.check_is_fitted()
+        return self.base_estimator_.predict(X)
+
+    def predict_proba(self, X):
+        """Return probability estimates for the test vector X.
+
+        Parameters
+        ----------
+        X : jax numpy array (n_samples, n_features) of object
+            Query points where the GP is evaluated for classification.
+
+        Returns
+        -------
+        C : jax numpy array (n_samples, n_classes)
+            Returns the probability of the samples for each class in
+            the model. The columns correspond to the classes in sorted
+            order.
+        """
+        self.check_is_fitted()
+        return self.base_estimator_.predict_proba(X)
+
+    def check_is_fitted(self):
+        """Perform is_fitted validation for estimator.
+
+        Checks if the estimator is fitted by verifying the presence of
+        fitted attribute self.n_classes_ and otherwise
+        raises a NotFittedError with the given message.
+
+        Raises
+        ------
+        Exception
+            If the attribute is not found.
+        """
+        try: 
+            self.n_classes_
+        except: 
+            raise Exception('Model is not fitted yet')
diff --git a/sml/gaussian_process/emulations/BUILD.bazel b/sml/gaussian_process/emulations/BUILD.bazel
index 7832e732..bf907a4f 100644
--- a/sml/gaussian_process/emulations/BUILD.bazel
+++ b/sml/gaussian_process/emulations/BUILD.bazel
@@ -11,3 +11,16 @@
 # 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_test(
+    name = "gpc_emul",
+    srcs = ["gpc_emul.py"],
+    deps = [
+        "//sml/gaussian_process:_gpc",
+        "//spu/utils:emulation",
+    ],
+)
diff --git a/sml/gaussian_process/emulations/gpc_emul.py b/sml/gaussian_process/emulations/gpc_emul.py
new file mode 100644
index 00000000..9e727d76
--- /dev/null
+++ b/sml/gaussian_process/emulations/gpc_emul.py
@@ -0,0 +1,49 @@
+import os
+import sys
+
+import jax.numpy as jnp
+from sklearn.datasets import load_iris
+
+# 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.gaussian_process._gpc import GaussianProcessClassifier
+
+
+def emul_gpc(mode: emulation.Mode.MULTIPROCESS):
+    def proc(x, y):
+        model = GaussianProcessClassifier(max_iter_predict=10)
+        model.fit(x, y)
+
+        pred = model.predict(x)
+        return pred
+
+    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_iris(return_X_y=True)
+        x = x[45:55, :]
+        y = y[45:55]
+
+        # mark these data to be protected in SPU
+        x, y = emulator.seal(x, y)
+        result = emulator.run(proc)(x, y)
+        print("Accuracy: ", jnp.sum(result == y) / len(y))
+
+    finally:
+        emulator.down()
+
+
+if __name__ == "__main__":
+    emul_gpc(emulation.Mode.MULTIPROCESS)
+
+    finally:
+        emulator.down()
+
+if __name__ == "__main__":
+    emul_gpc(emulation.Mode.MULTIPROCESS)
diff --git a/sml/gaussian_process/kernels.py b/sml/gaussian_process/kernels.py
new file mode 100644
index 00000000..40e9c59c
--- /dev/null
+++ b/sml/gaussian_process/kernels.py
@@ -0,0 +1,15 @@
+import jax.numpy as jnp
+
+
+class RBF:
+    def __init__(self, length_scale=1.0):
+        self.length_scale = length_scale
+
+    def __call__(self, X, Y=None):
+        if Y == None:
+            Y = X
+        K = jnp.zeros((X.shape[0], Y.shape[0]), dtype=jnp.float32)
+        for i in range(X.shape[0]):
+            for j in range(Y.shape[0]):
+                K = K.at[i, j].set((jnp.sum((X[i] - Y[j]) ** 2)))
+        return jnp.exp(-K / (2 * self.length_scale))
diff --git a/sml/gaussian_process/ovo_ovr.py b/sml/gaussian_process/ovo_ovr.py
new file mode 100644
index 00000000..d7c99b5a
--- /dev/null
+++ b/sml/gaussian_process/ovo_ovr.py
@@ -0,0 +1,109 @@
+import jax.numpy as jnp
+from jax import vmap
+from jax.scipy.special import erf, expit
+from kernels import RBF
+from jax.lax.linalg import cholesky
+from jax.scipy.linalg import cho_solve, solve
+
+LAMBDAS = jnp.array([0.41, 0.4, 0.37, 0.44, 0.39])[:, jnp.newaxis]
+COEFS = jnp.array(
+    [-1854.8214151, 3516.89893646, 221.29346712, 128.12323805, -2010.49422654]
+)[:, jnp.newaxis]
+
+class OneVsRestClassifier:
+    def __init__(self, estimator, n_classes):
+        self.estimator = estimator
+        self.classes_ = n_classes
+
+    def fit(self, X, y):
+
+        self.estimator.approx_func = expit
+        self.estimator.X_train_ = jnp.array(X)
+
+        if self.estimator.kernel is None:  # Use an RBF kernel as default
+            self.estimator.kernel_ = RBF()
+        else:
+            self.estimator.kernel_ = self.estimator.kernel
+
+        self.K = self.estimator.kernel_(X)
+
+        self.y_binary = jnp.array(
+            [jnp.where(y == i, 0, 1) for i in range(self.classes_)]
+        )
+
+        self.fs_ = vmap(self._posterior_mode, in_axes=(None, 0, None))(X, self.y_binary, self.K)
+        self.f_cached = self.fs_
+
+    def predict(self, X_test):
+        X = jnp.array(X_test)
+        K_star = self.estimator.kernel_(self.estimator.X_train_, X)
+        diag_K_Xtest = jnp.diag(self.estimator.kernel_(X))
+        maxima = vmap(self.predict_proba_oneclass, in_axes=(0, None, 0, None, None, None))(
+            self.y_binary, X_test, self.fs_, self.K, K_star, diag_K_Xtest
+        )
+        return maxima.argmax(axis=0)
+
+    def predict_proba(self, X_test):
+        X = jnp.array(X_test)
+        K_star = self.estimator.kernel_(self.estimator.X_train_, X)
+        diag_K_Xtest = jnp.diag(self.estimator.kernel_(X))
+        maxima = vmap(self.predict_proba_oneclass, in_axes=(0, None, 0, None, None, None))(
+            self.y_binary, X_test, self.fs_, self.K, K_star, diag_K_Xtest
+        )
+        maxima = maxima / jnp.sum(maxima, axis=0)
+        return maxima.T
+
+    def predict_proba_oneclass(self, y_binary, X_test, f_, K, K_star, diag_K_Xtest):
+        X = jnp.asarray(X_test)
+        # K = self.estimator.kernel_(self.estimator.X_train_)
+
+        pi = self.estimator.approx_func(f_)
+        W = pi * (1 - pi)
+
+        W_sqr = jnp.sqrt(W)
+        W_sqr_K = W_sqr[:, jnp.newaxis] * K
+        B = jnp.eye(W.shape[0]) + W_sqr_K * W_sqr
+        L = cholesky(B)
+
+        # K_star = self.estimator.kernel_(self.estimator.X_train_, X)
+        # f_star = K_star.T.dot(self.log_and_grad(self.f_, self.y_train))
+        f_star = K_star.T.dot(y_binary - pi)
+        v = solve(L, W_sqr[:, jnp.newaxis] * K_star)
+        var_f_star = diag_K_Xtest - jnp.einsum("ij,ij->j", v, v)
+
+        alpha = 1 / (2 * var_f_star)
+        gamma = LAMBDAS * f_star
+        integrals = (
+            jnp.sqrt(jnp.pi / alpha)
+            * erf(gamma * jnp.sqrt(alpha / (alpha + LAMBDAS**2)))
+            / (2 * jnp.sqrt(var_f_star * 2 * jnp.pi))
+        )
+        pi_star = (COEFS * integrals).sum(axis=0) + 0.5 * COEFS.sum()
+
+        return 1 - pi_star
+
+    def _posterior_mode(self, X, y_binary, K):
+        # K = self.estimator.kernel_(X)
+        # Based on Algorithm 3.1 of GPML
+        f = jnp.zeros_like(
+            y_binary, dtype=jnp.float32
+        )  # a warning is triggered if float64 is used
+
+        for _ in range(self.estimator.max_iter_predict):
+            # W = self.log_and_2grads_and_negtive(f, self.y_train)
+            pi =  self.estimator.approx_func(f)
+            W = pi * (1 - pi)
+            W_sqr = jnp.sqrt(W)
+            W_sqr_K = W_sqr[:, jnp.newaxis] * K
+
+            B = jnp.eye(W.shape[0]) + W_sqr_K * W_sqr
+            L = cholesky(B)
+            # b = W * f + self.log_and_grad(f, self.y_train)
+            b = W * f + (y_binary - pi)
+            a = b - jnp.dot(
+                W_sqr[:, jnp.newaxis] * cho_solve((L, True), jnp.eye(W.shape[0])),
+                W_sqr_K.dot(b),
+            )
+            f = K.dot(a)
+
+        return f


Test patch:
diff --git a/sml/gaussian_process/tests/BUILD.bazel b/sml/gaussian_process/tests/BUILD.bazel
index 7832e732..52e9a0a3 100644
--- a/sml/gaussian_process/tests/BUILD.bazel
+++ b/sml/gaussian_process/tests/BUILD.bazel
@@ -11,3 +11,17 @@
 # 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 = "gpc_test",
+    srcs = ["gpc_test.py"],
+    deps = [
+        "//sml/gaussian_process:_gpc",
+        "//spu:init",
+        "//spu/utils:simulation",
+    ],
+)
diff --git a/sml/gaussian_process/tests/gpc_test.py b/sml/gaussian_process/tests/gpc_test.py
new file mode 100644
index 00000000..a854252b
--- /dev/null
+++ b/sml/gaussian_process/tests/gpc_test.py
@@ -0,0 +1,43 @@
+import os
+import sys
+import unittest
+
+import jax
+import jax.numpy as jnp
+from sklearn.datasets import load_iris
+
+import spu.spu_pb2 as spu_pb2
+import spu.utils.simulation as spsim
+
+# Add the library directory to the path
+sys.path.append(os.path.join(os.path.dirname(__file__), "../../../"))
+from sml.gaussian_process._gpc import GaussianProcessClassifier
+
+
+class UnitTests(unittest.TestCase):
+    def test_gpc(self):
+        sim = spsim.Simulator.simple(
+            3, spu_pb2.ProtocolKind.ABY3, spu_pb2.FieldType.FM64
+        )
+
+        # Test GaussianProcessClassifier
+        def proc(x, y):
+            model = GaussianProcessClassifier(max_iter_predict=10)
+            model.fit(x, y, 3)
+
+            pred = model.predict(x)
+            return pred
+
+        # Create dataset
+        x, y = load_iris(return_X_y=True)
+        # x = x[45:55, :]
+        # y = y[45:55]
+
+        # Run
+        result = spsim.sim_jax(sim, proc)(x, y)
+        print(result)
+        print(y)
+        print("Accuracy: ", jnp.sum(result == y) / len(y))
+
+
+unittest.main()