grouped_fe.cluster_embeddings

Guide & worked example · Source

grouped_fe.cluster_embeddings(
    embeddings,
    n_groups,
    niter=25,
    tol=1e-06,
    seed=0,
    use_flash_kmeans=True,
)

Cluster embeddings with flash-kmeans when available, otherwise torch fallback.