Update classification_test.py

secretflow · Nov 17, 2023 · 90e3cd2 · 90e3cd2
1 parent c1decdb
commit 90e3cd2
Showing 1 changed file with 2 additions and 4 deletions.
diff --git a/sml/metrics/classification/classification_test.py b/sml/metrics/classification/classification_test.py
@@ -139,16 +139,14 @@ def check(spu_result, sk_result):
         # Test binary
         y_true = jnp.array([0, 1, 1, 0, 1, 1])
         y_pred = jnp.array([0, 0, 1, 0, 1, 1])
-        spu_result = spsim.sim_jax(sim, proc)(y_true, y_pred, pos_label=1, transform=0)
+        spu_result = spsim.sim_jax(sim, proc, static_argnums=(2, 5))(y_true, y_pred, 'binary', None, 1, False)
         sk_result = sklearn_proc(y_true, y_pred)
         check(spu_result, sk_result)
 
         # Test multiclass
         y_true = jnp.array([0, 1, 1, 0, 2, 1])
         y_pred = jnp.array([0, 0, 1, 0, 2, 1])
-        spu_result = spsim.sim_jax(sim, proc)(
-            y_true, y_pred, average=None, labels=[0, 1, 2]
-        )
+        spu_result = spsim.sim_jax(sim, proc, static_argnums=(2, 5))(y_true, y_pred, None, [0,1,2], 1, True)
         sk_result = sklearn_proc(y_true, y_pred, average=None, labels=[0, 1, 2])
         check(spu_result, sk_result)