spu: [Bug]: The behavior of kmeans in SPU does not match kmeans in sklearns

Issue Type

Currentness/Accuracy

Modules Involved

Others

Have you reproduced the bug with SPU HEAD?

Yes

Have you searched existing issues?

Yes

SPU Version

spu 0.7.0b0

OS Platform and Distribution

Linux Ubuntu 22.04

Python Version

3.10

Compiler Version

No response

Current Behavior?

Hello, sorry for disturbing you.

When I use Kmeans implemented in sml/cluster, I encounter the problem that the behavior of it does not match kmeans implemented in sklearns with some inputs. I will include the code below.

Standalone code to reproduce the issue

sim = spsim.Simulator.simple(
    3, spu_pb2.ProtocolKind.ABY3, spu_pb2.FieldType.FM64
)
def proc(x):
    model = KMEANS(n_clusters=3, n_samples=x.shape[0], max_iter=10)
    model.fit(x)
    return model._centers
X = jnp.array([[-4, -3, -2, -1]]).T
result = spsim.sim_jax(sim, proc)(X)
print("result\n", result)

# Compare with sklearn
from sklearn.cluster import KMeans

model = KMeans(n_clusters=3)
model.fit(X)
print("sklearn:\n", model.cluster_centers_)

Relevant log output

result
 [[-1.4999962]
 [ 0.       ]
 [-3.4999924]]
sklearn:
 [[-1. ]
 [-3.5]
 [-2. ]]