def cost(points, centers):
a_sqr = jnp.sum(points * points, 1)[None, :]
b_sqr = jnp.sum(centers * centers, 1)[:, None]
diff = jnp.matmul(points, centers.T).T
dists = sparse.todense(a_sqr) + sparse.todense(
b_sqr) - 2 * sparse.todense(diff)
min_dist = jnp.min(dists, axis=0)
return min_dist.sum()
def func(M):
return todense(M) + 1