diff --git a/dynamax/hidden_markov_model/models/arhmm.py b/dynamax/hidden_markov_model/models/arhmm.py
index 222f4f31..c8961a04 100644
--- a/dynamax/hidden_markov_model/models/arhmm.py
+++ b/dynamax/hidden_markov_model/models/arhmm.py
@@ -15,6 +15,7 @@
 from dynamax.parameters import ParameterProperties
 from dynamax.types import Scalar
 from dynamax.utils.bijectors import RealToPSDBijector
+from dynamax.utils.cluster import kmeans
 from tensorflow_probability.substrates import jax as tfp
 
 tfd = tfp.distributions
@@ -64,12 +65,8 @@ def initialize(self,
 
         if method.lower() == "kmeans":
             assert emissions is not None, "Need emissions to initialize the model with K-Means!"
-            from sklearn.cluster import KMeans
-            key, subkey = jr.split(key)  # Create a random seed for SKLearn.
-            sklearn_key = jr.randint(subkey, shape=(), minval=0, maxval=2147483647)  # Max int32 value.
-            km = KMeans(self.num_states, random_state=int(sklearn_key)).fit(emissions.reshape(-1, self.emission_dim))
             _emission_weights = jnp.zeros((self.num_states, self.emission_dim, self.emission_dim * self.num_lags))
-            _emission_biases = jnp.array(km.cluster_centers_)
+            _emission_biases = kmeans(emissions.reshape(-1, self.emission_dim), self.num_states, key).centroids
             _emission_covs = jnp.tile(jnp.eye(self.emission_dim)[None, :, :], (self.num_states, 1, 1))
 
         elif method.lower() == "prior":
diff --git a/dynamax/hidden_markov_model/models/gamma_hmm.py b/dynamax/hidden_markov_model/models/gamma_hmm.py
index 6bf165fd..20fb8088 100644
--- a/dynamax/hidden_markov_model/models/gamma_hmm.py
+++ b/dynamax/hidden_markov_model/models/gamma_hmm.py
@@ -9,6 +9,7 @@
 from dynamax.hidden_markov_model.models.initial import StandardHMMInitialState, ParamsStandardHMMInitialState
 from dynamax.hidden_markov_model.models.transitions import StandardHMMTransitions, ParamsStandardHMMTransitions
 from dynamax.types import Scalar
+from dynamax.utils.cluster import kmeans
 import optax
 from typing import NamedTuple, Optional, Tuple, Union
 
@@ -67,14 +68,8 @@ def initialize(
 
         if method.lower() == "kmeans":
             assert emissions is not None, "Need emissions to initialize the model with K-Means!"
-            from sklearn.cluster import KMeans
-
-            key, subkey = jr.split(key)  # Create a random seed for SKLearn.
-            sklearn_key = jr.randint(subkey, shape=(), minval=0, maxval=2147483647)  # Max int32 value.
-            km = KMeans(self.num_states, random_state=int(sklearn_key)).fit(emissions.reshape(-1, 1))
-
             _emission_concentrations = jnp.ones((self.num_states,))
-            _emission_rates = jnp.ravel(1.0 / km.cluster_centers_)
+            _emission_rates = jnp.ravel(1.0 / kmeans(emissions.reshape(-1, 1), self.num_states, key).centroids)
 
         elif method.lower() == "prior":
             _emission_concentrations = jnp.ones((self.num_states,))
diff --git a/dynamax/hidden_markov_model/models/gaussian_hmm.py b/dynamax/hidden_markov_model/models/gaussian_hmm.py
index a959188b..74f7630d 100644
--- a/dynamax/hidden_markov_model/models/gaussian_hmm.py
+++ b/dynamax/hidden_markov_model/models/gaussian_hmm.py
@@ -21,6 +21,7 @@
 from dynamax.utils.distributions import niw_posterior_update
 from dynamax.utils.bijectors import RealToPSDBijector
 from dynamax.utils.utils import pytree_sum
+from dynamax.utils.cluster import kmeans
 
 
 class ParamsGaussianHMMEmissions(NamedTuple):
@@ -97,12 +98,7 @@ def initialize(self,
         """
         if method.lower() == "kmeans":
             assert emissions is not None, "Need emissions to initialize the model with K-Means!"
-            from sklearn.cluster import KMeans
-            key, subkey = jr.split(key)  # Create a random seed for SKLearn.
-            sklearn_key = jr.randint(subkey, shape=(), minval=0, maxval=2147483647)  # Max int32 value.
-            km = KMeans(self.num_states, random_state=int(sklearn_key)).fit(emissions.reshape(-1, self.emission_dim))
-
-            _emission_means = jnp.array(km.cluster_centers_)
+            _emission_means = kmeans(emissions.reshape(-1, self.emission_dim), self.num_states, key).centroids
             _emission_covs = jnp.tile(jnp.eye(self.emission_dim)[None, :, :], (self.num_states, 1, 1))
 
         elif method.lower() == "prior":
@@ -239,11 +235,7 @@ def initialize(self,
 
         if method.lower() == "kmeans":
             assert emissions is not None, "Need emissions to initialize the model with K-Means!"
-            from sklearn.cluster import KMeans
-            key, subkey = jr.split(key)  # Create a random seed for SKLearn.
-            sklearn_key = jr.randint(subkey, shape=(), minval=0, maxval=2147483647)  # Max int32 value.
-            km = KMeans(self.num_states, random_state=int(sklearn_key)).fit(emissions.reshape(-1, self.emission_dim))
-            _emission_means = jnp.array(km.cluster_centers_)
+            _emission_means = kmeans(emissions.reshape(-1, self.emission_dim), self.num_states, key).centroids
             _emission_scale_diags = jnp.ones((self.num_states, self.emission_dim))
 
         elif method.lower() == "prior":
@@ -383,11 +375,7 @@ def initialize(self,
         """
         if method.lower() == "kmeans":
             assert emissions is not None, "Need emissions to initialize the model with K-Means!"
-            from sklearn.cluster import KMeans
-            key, subkey = jr.split(key)  # Create a random seed for SKLearn.
-            sklearn_key = jr.randint(subkey, shape=(), minval=0, maxval=2147483647)  # Max int32 value.
-            km = KMeans(self.num_states, random_state=int(sklearn_key)).fit(emissions.reshape(-1, self.emission_dim))
-            _emission_means = jnp.array(km.cluster_centers_)
+            _emission_means = kmeans(emissions.reshape(-1, self.emission_dim), self.num_states, key).centroids
             _emission_scales = jnp.ones((self.num_states,))
 
         elif method.lower() == "prior":
@@ -498,11 +486,7 @@ def initialize(self,
         """
         if method.lower() == "kmeans":
             assert emissions is not None, "Need emissions to initialize the model with K-Means!"
-            from sklearn.cluster import KMeans
-            key, subkey = jr.split(key)  # Create a random seed for SKLearn.
-            sklearn_key = jr.randint(subkey, shape=(), minval=0, maxval=2147483647)  # Max int32 value.
-            km = KMeans(self.num_states, random_state=int(sklearn_key)).fit(emissions.reshape(-1, self.emission_dim))
-            _emission_means = jnp.array(km.cluster_centers_)
+            _emission_means = kmeans(emissions.reshape(-1, self.emission_dim), self.num_states, key).centroids
             _emission_cov = jnp.eye(self.emission_dim)
 
         elif method.lower() == "prior":
@@ -663,11 +647,7 @@ def initialize(self,
         """
         if method.lower() == "kmeans":
             assert emissions is not None, "Need emissions to initialize the model with K-Means!"
-            from sklearn.cluster import KMeans
-            key, subkey = jr.split(key)  # Create a random seed for SKLearn.
-            sklearn_key = jr.randint(subkey, shape=(), minval=0, maxval=2147483647)  # Max int32 value.
-            km = KMeans(self.num_states, random_state=int(sklearn_key)).fit(emissions.reshape(-1, self.emission_dim))
-            _emission_means = jnp.array(km.cluster_centers_)
+            _emission_means = kmeans(emissions.reshape(-1, self.emission_dim), self.num_states, key).centroids
             _emission_cov_diag_factors = jnp.ones((self.num_states, self.emission_dim))
             _emission_cov_low_rank_factors = jnp.zeros((self.num_states, self.emission_dim, self.emission_rank))
 
diff --git a/dynamax/hidden_markov_model/models/gmm_hmm.py b/dynamax/hidden_markov_model/models/gmm_hmm.py
index c4646b02..d6dc346b 100644
--- a/dynamax/hidden_markov_model/models/gmm_hmm.py
+++ b/dynamax/hidden_markov_model/models/gmm_hmm.py
@@ -20,6 +20,7 @@
 from dynamax.hidden_markov_model.models.transitions import StandardHMMTransitions, ParamsStandardHMMTransitions
 from dynamax.utils.bijectors import RealToPSDBijector
 from dynamax.utils.utils import pytree_sum
+from dynamax.utils.cluster import kmeans
 from dynamax.types import IntScalar, Scalar
 
 
@@ -114,12 +115,11 @@ def initialize(self,
         """
         if method.lower() == "kmeans":
             assert emissions is not None, "Need emissions to initialize the model with K-Means!"
-            from sklearn.cluster import KMeans
-            key, subkey = jr.split(key)  # Create a random seed for SKLearn.
-            sklearn_key = jr.randint(subkey, shape=(), minval=0, maxval=2147483647)  # Max int32 value.
-            km = KMeans(self.num_states, random_state=int(sklearn_key)).fit(emissions.reshape(-1, self.emission_dim))
             _emission_weights = jnp.ones((self.num_states, self.num_components)) / self.num_components
-            _emission_means = jnp.tile(jnp.array(km.cluster_centers_)[:, None, :], (1, self.num_components, 1))
+            _emission_means = jnp.tile(
+                kmeans(emissions.reshape(-1, self.emission_dim), self.num_states, key).centroids[:, None, :],
+                (1, self.num_components, 1),
+            )
             _emission_covs = jnp.tile(jnp.eye(self.emission_dim), (self.num_states, self.num_components, 1, 1))
 
         elif method.lower() == "prior":
@@ -394,12 +394,11 @@ def initialize(self,
         """
         if method.lower() == "kmeans":
             assert emissions is not None, "Need emissions to initialize the model with K-Means!"
-            from sklearn.cluster import KMeans
-            key, subkey = jr.split(key)  # Create a random seed for SKLearn.
-            sklearn_key = jr.randint(subkey, shape=(), minval=0, maxval=2147483647)  # Max int32 value.
-            km = KMeans(self.num_states, random_state=int(sklearn_key)).fit(emissions.reshape(-1, self.emission_dim))
             _emission_weights = jnp.ones((self.num_states, self.num_components)) / self.num_components
-            _emission_means = jnp.tile(jnp.array(km.cluster_centers_)[:, None, :], (1, self.num_components, 1))
+            _emission_means = jnp.tile(
+                kmeans(emissions.reshape(-1, self.emission_dim), self.num_states, key).centroids[:, None, :],
+                (1, self.num_components, 1),
+            )
             _emission_scale_diags = jnp.ones((self.num_states, self.num_components, self.emission_dim))
 
         elif method.lower() == "prior":
diff --git a/dynamax/hidden_markov_model/models/linreg_hmm.py b/dynamax/hidden_markov_model/models/linreg_hmm.py
index f6bc2900..3838de72 100644
--- a/dynamax/hidden_markov_model/models/linreg_hmm.py
+++ b/dynamax/hidden_markov_model/models/linreg_hmm.py
@@ -16,6 +16,7 @@
 from dynamax.types import Scalar
 from dynamax.utils.utils import pytree_sum
 from dynamax.utils.bijectors import RealToPSDBijector
+from dynamax.utils.cluster import kmeans
 
 tfd = tfp.distributions
 tfb = tfp.bijectors
@@ -80,12 +81,8 @@ def initialize(self,
         """
         if method.lower() == "kmeans":
             assert emissions is not None, "Need emissions to initialize the model with K-Means!"
-            from sklearn.cluster import KMeans
-            key, subkey = jr.split(key)  # Create a random seed for SKLearn.
-            sklearn_key = jr.randint(subkey, shape=(), minval=0, maxval=2147483647)  # Max int32 value.
-            km = KMeans(self.num_states, random_state=int(sklearn_key)).fit(emissions.reshape(-1, self.emission_dim))
             _emission_weights = jnp.zeros((self.num_states, self.emission_dim, self.input_dim))
-            _emission_biases = jnp.array(km.cluster_centers_)
+            _emission_biases = kmeans(emissions.reshape(-1, self.emission_dim), self.num_states, key).centroids
             _emission_covs = jnp.tile(jnp.eye(self.emission_dim)[None, :, :], (self.num_states, 1, 1))
 
         elif method.lower() == "prior":
diff --git a/dynamax/hidden_markov_model/models/logreg_hmm.py b/dynamax/hidden_markov_model/models/logreg_hmm.py
index e2199fff..caec1736 100644
--- a/dynamax/hidden_markov_model/models/logreg_hmm.py
+++ b/dynamax/hidden_markov_model/models/logreg_hmm.py
@@ -14,6 +14,7 @@
 from dynamax.hidden_markov_model.models.initial import StandardHMMInitialState, ParamsStandardHMMInitialState
 from dynamax.hidden_markov_model.models.transitions import StandardHMMTransitions, ParamsStandardHMMTransitions
 from dynamax.types import IntScalar, Scalar
+from dynamax.utils.cluster import kmeans
 
 
 class ParamsLogisticRegressionHMMEmissions(NamedTuple):
@@ -80,16 +81,19 @@ def initialize(self,
         if method.lower() == "kmeans":
             assert emissions is not None, "Need emissions to initialize the model with K-Means!"
             assert inputs is not None, "Need inputs to initialize the model with K-Means!"
-            from sklearn.cluster import KMeans
-
             flat_emissions = emissions.reshape(-1,)
             flat_inputs = inputs.reshape(-1, self.input_dim)
-            key, subkey = jr.split(key)  # Create a random seed for SKLearn.
-            sklearn_key = jr.randint(subkey, shape=(), minval=0, maxval=2147483647)  # Max int32 value.
-            km = KMeans(self.num_states, random_state=int(sklearn_key)).fit(flat_inputs)
+
+            assignments = kmeans(flat_inputs, self.num_states, key).assignments
             _emission_weights = jnp.zeros((self.num_states, self.input_dim))
-            _emission_biases = jnp.array([tfb.Sigmoid().inverse(flat_emissions[km.labels_ == k].mean())
-                                          for k in range(self.num_states)])
+            # Both guards keep the logit below finite: an empty cluster has an undefined
+            # mean, and an all-0s or all-1s cluster would otherwise logit to -/+ infinity.
+            cluster_means = jnp.array(
+                [jnp.mean(flat_emissions, where=(assignments == k)) for k in range(self.num_states)]
+            )
+            cluster_means = jnp.where(jnp.isnan(cluster_means), flat_emissions.mean(), cluster_means)
+            eps = jnp.finfo(cluster_means.dtype).eps
+            _emission_biases = tfb.Sigmoid().inverse(jnp.clip(cluster_means, eps, 1.0 - eps))
 
         elif method.lower() == "prior":
             # TODO: Use an MNIW prior
diff --git a/dynamax/hidden_markov_model/models/test_models.py b/dynamax/hidden_markov_model/models/test_models.py
index 73a2969a..e177b2b5 100644
--- a/dynamax/hidden_markov_model/models/test_models.py
+++ b/dynamax/hidden_markov_model/models/test_models.py
@@ -1,6 +1,7 @@
 """Tests for the HMM models."""
 
 import dynamax.hidden_markov_model as models
+import jax
 import jax.numpy as jnp
 import jax.random as jr
 import pytest
@@ -169,3 +170,127 @@ def test_sample_and_fit_arhmm():
 #     test_cov = test_hmm.emission_covariance_matrices.value
 #     assert jnp.alltrue(test_cov.shape == (10, 2, 2))
 #     assert jnp.allclose(jnp.linalg.norm(test_cov-refr_cov, axis=-1), 0., atol=1)
+
+
+# Models whose emissions support kmeans initialization. BernoulliHMM, CategoricalHMM,
+# MultinomialHMM and PoissonHMM deliberately raise NotImplementedError instead.
+KMEANS_CONFIGS = [
+    # (cls, kwargs, inputs, init_takes_inputs)
+    (models.GammaHMM, dict(num_states=4), None, False),
+    (models.GaussianHMM, dict(num_states=4, emission_dim=3), None, False),
+    (models.DiagonalGaussianHMM, dict(num_states=4, emission_dim=3), None, False),
+    (models.SphericalGaussianHMM, dict(num_states=4, emission_dim=3), None, False),
+    (models.SharedCovarianceGaussianHMM, dict(num_states=4, emission_dim=3), None, False),
+    (models.LowRankGaussianHMM, dict(num_states=4, emission_dim=3, emission_rank=1), None, False),
+    (models.GaussianMixtureHMM, dict(num_states=4, num_components=2, emission_dim=3), None, False),
+    (models.DiagonalGaussianMixtureHMM, dict(num_states=4, num_components=2, emission_dim=3), None, False),
+    (models.LinearRegressionHMM, dict(num_states=3, emission_dim=3, input_dim=5),
+     jr.normal(jr.PRNGKey(0), (NUM_TIMESTEPS, 5)), False),
+    (models.LogisticRegressionHMM, dict(num_states=4, input_dim=5),
+     jr.normal(jr.PRNGKey(0), (NUM_TIMESTEPS, 5)), True),
+]
+
+
+@pytest.mark.parametrize(["cls", "kwargs", "inputs", "init_takes_inputs"], KMEANS_CONFIGS)
+def test_initialize_kmeans(cls, kwargs, inputs, init_takes_inputs):
+    """Test that kmeans initialization produces finite, fittable parameters."""
+    hmm = cls(**kwargs)
+    key1, key2, key3 = jr.split(jr.PRNGKey(42), 3)
+    params, _ = hmm.initialize(key1)
+    _, emissions = hmm.sample(params, key2, num_timesteps=NUM_TIMESTEPS, inputs=inputs)
+
+    # Only LogisticRegressionHMM clusters its inputs, so only it accepts them here.
+    init_kwargs = dict(emissions=emissions)
+    if init_takes_inputs:
+        init_kwargs["inputs"] = inputs
+    km_params, km_props = hmm.initialize(key3, method="kmeans", **init_kwargs)
+
+    # No parameter may come back NaN or infinite; that was the failure mode of an
+    # unguarded empty cluster.
+    for leaf in jax.tree_util.tree_leaves(km_params):
+        assert jnp.all(jnp.isfinite(leaf))
+
+    # The initialization must be usable: EM from it improves monotonically.
+    _, lps = hmm.fit_em(km_params, km_props, emissions, inputs=inputs, num_iters=3, verbose=False)
+    assert monotonically_increasing(lps, atol=1e-2, rtol=1e-2)
+
+
+def test_initialize_kmeans_arhmm():
+    """Test that kmeans initialization works for a LinearAutoregressiveHMM."""
+    arhmm = models.LinearAutoregressiveHMM(num_states=4, emission_dim=2, num_lags=1)
+    key1, key2, key3 = jr.split(jr.PRNGKey(42), 3)
+    params, _ = arhmm.initialize(key1)
+    _, emissions = arhmm.sample(params, key2, num_timesteps=NUM_TIMESTEPS)
+
+    km_params, _ = arhmm.initialize(key3, method="kmeans", emissions=emissions)
+
+    for leaf in jax.tree_util.tree_leaves(km_params):
+        assert jnp.all(jnp.isfinite(leaf))
+    assert km_params.emissions.biases.shape == (4, 2)
+
+
+def test_logreg_hmm_kmeans_finite_bias_with_saturated_cluster():
+    """A cluster whose assigned emissions are all 0 (or all 1) is not the same as an
+    empty cluster, but it drives the cluster mean to exactly 0.0 or 1.0. Feeding that
+    straight into the logit used to produce +/-inf biases even though the existing
+    NaN guard (for genuinely empty clusters) passed. Construct inputs as two
+    well-separated blobs so kmeans reliably keeps them apart, and make one blob's
+    emissions uniformly 0 and the other uniformly 1, so the cluster means saturate
+    at both 0.0 and 1.0 and exercise both the lower and upper clip bounds.
+    """
+    from dynamax.hidden_markov_model.models.logreg_hmm import LogisticRegressionHMMEmissions
+
+    emission_component = LogisticRegressionHMMEmissions(num_states=2, input_dim=2)
+    inputs = jnp.concatenate([
+        jnp.tile(jnp.array([1000.0, 0.0]), (10, 1)),
+        jnp.tile(jnp.array([-1000.0, 0.0]), (10, 1)),
+    ], axis=0)
+    emissions = jnp.concatenate([
+        jnp.zeros(10),
+        jnp.ones(10),
+    ], axis=0)
+
+    for seed in range(10):
+        params, _ = emission_component.initialize(
+            jr.PRNGKey(seed), method="kmeans", emissions=emissions, inputs=inputs)
+        assert jnp.all(jnp.isfinite(params.biases)), f"seed {seed} produced a non-finite bias"
+        # Sanity check that this seed really did hit both the low- and high-saturated
+        # cluster cases (the clip's eps floor and its 1 - eps ceiling).
+        assert jnp.any(params.biases < 0)
+        assert jnp.any(params.biases > 0)
+
+
+def test_initialize_kmeans_is_jax_transformable():
+    """Kmeans initialization must stay compatible with `jax.jit` and `jax.vmap`.
+
+    `LogisticRegressionHMMEmissions.initialize`'s kmeans branch computes each
+    cluster's emission mean with `jnp.mean(flat_emissions, where=(assignments == k))`.
+    The seemingly-equivalent `flat_emissions[assignments == k].mean()` also passes
+    eagerly, but it boolean-mask-indexes with a *traced* array, which produces a
+    variable-shaped intermediate. That is illegal under `jax.jit`/`jax.vmap` (it
+    raises `NonConcreteBooleanIndexError`), even though nothing catches it outside
+    a JAX transformation. If this test starts failing with that error, the fix is
+    to restore the `where=` form in the kmeans branch of `initialize` -- not to
+    delete this test.
+    """
+    num_states, input_dim = 4, 5
+    hmm = models.LogisticRegressionHMM(num_states=num_states, input_dim=input_dim)
+    key1, key2 = jr.split(jr.PRNGKey(0), 2)
+    params, _ = hmm.initialize(key1)
+    inputs = jr.normal(key2, (NUM_TIMESTEPS, input_dim))
+    _, emissions = hmm.sample(params, jr.PRNGKey(1), num_timesteps=NUM_TIMESTEPS, inputs=inputs)
+
+    def init_kmeans(key, emissions, inputs):
+        km_params, _ = hmm.initialize(key, method="kmeans", emissions=emissions, inputs=inputs)
+        return km_params
+
+    jitted_params = jax.jit(init_kmeans)(jr.PRNGKey(2), emissions, inputs)
+    assert jnp.all(jnp.isfinite(jitted_params.emissions.biases))
+
+    batch_size = 3
+    batch_keys = jr.split(jr.PRNGKey(3), batch_size)
+    batch_emissions = jnp.stack([emissions] * batch_size)
+    batch_inputs = jnp.stack([inputs] * batch_size)
+    vmapped_params = vmap(init_kmeans)(batch_keys, batch_emissions, batch_inputs)
+    assert vmapped_params.emissions.biases.shape == (batch_size, num_states)
+    assert jnp.all(jnp.isfinite(vmapped_params.emissions.biases))
diff --git a/dynamax/utils/cluster.py b/dynamax/utils/cluster.py
new file mode 100644
index 00000000..998d3507
--- /dev/null
+++ b/dynamax/utils/cluster.py
@@ -0,0 +1,227 @@
+"""K-means clustering in JAX, used to initialize HMM emission parameters."""
+
+import math
+from functools import partial
+from typing import NamedTuple, Optional
+
+from jax import jit, lax
+from jax import numpy as jnp
+from jax import random as jr
+from jaxtyping import Array, Float, Int
+
+from dynamax.types import PRNGKeyT, Scalar
+
+
+class KMeansState(NamedTuple):
+    """Result of a k-means fit.
+
+    Attributes:
+        centroids: cluster centers.
+        assignments: index of the closest centroid for each sample.
+        inertia: sum of squared distances from each sample to its centroid.
+        n_iter: number of Lloyd iterations run by the selected restart.
+    """
+
+    centroids: Float[Array, "num_clusters num_features"]
+    assignments: Int[Array, " num_samples"]
+    inertia: Float[Array, ""]
+    n_iter: Int[Array, ""]
+
+
+def _squared_distances(
+    X: Float[Array, "num_samples num_features"],
+    centroids: Float[Array, "num_clusters num_features"],
+) -> Float[Array, "num_samples num_clusters"]:
+    """Compute squared euclidean distances from every sample to every centroid.
+
+    Expands ||x - c||^2 to ||x||^2 - 2 x.c + ||c||^2 so that no
+    (num_samples, num_clusters, num_features) intermediate is materialized.
+    Centers both operands on the sample mean first: the uncentered expansion
+    is not offset-stable in float32 and returns negative "squared" distances
+    once the data sits a few orders of magnitude from the origin, which
+    silently corrupts both clustering and n_init restart selection.
+    """
+    offset = jnp.mean(X, axis=0)
+    centered_X = X - offset
+    centered_centroids = centroids - offset
+    return jnp.maximum(
+        jnp.sum(centered_X**2, axis=1)[:, None]
+        - 2.0 * centered_X @ centered_centroids.T
+        + jnp.sum(centered_centroids**2, axis=1)[None, :],
+        0.0,
+    )
+
+
+def _assign(
+    X: Float[Array, "num_samples num_features"],
+    centroids: Float[Array, "num_clusters num_features"],
+) -> Int[Array, " num_samples"]:
+    """Assign each sample to its closest centroid."""
+    return jnp.argmin(_squared_distances(X, centroids), axis=1)
+
+
+def _update_centroids(
+    X: Float[Array, "num_samples num_features"],
+    assignments: Int[Array, " num_samples"],
+    num_clusters: int,
+    previous: Float[Array, "num_clusters num_features"],
+) -> Float[Array, "num_clusters num_features"]:
+    """Recompute centroids as the mean of their assigned samples.
+
+    A cluster that captured no samples retains its previous centroid. Averaging an
+    empty cluster would produce NaN, which propagates into the emission parameters
+    and also prevents the fixed-point loop from ever terminating.
+    """
+    num_features = X.shape[1]
+    sums = jnp.zeros((num_clusters, num_features), X.dtype).at[assignments].add(X)
+    counts = jnp.zeros((num_clusters,), X.dtype).at[assignments].add(jnp.ones((), X.dtype))
+    means = sums / jnp.maximum(counts, 1.0)[:, None]
+    return jnp.where(counts[:, None] > 0, means, previous)
+
+
+def _inertia(
+    X: Float[Array, "num_samples num_features"],
+    centroids: Float[Array, "num_clusters num_features"],
+) -> Float[Array, ""]:
+    """Sum of squared distances from each sample to its closest centroid."""
+    return jnp.sum(jnp.min(_squared_distances(X, centroids), axis=1))
+
+
+def _kmeans_plusplus(
+    key: PRNGKeyT,
+    X: Float[Array, "num_samples num_features"],
+    num_clusters: int,
+    n_local_trials: int,
+) -> Float[Array, "num_clusters num_features"]:
+    """Choose initial centroids with greedy k-means++.
+
+    Samples each successive centroid with probability proportional to its squared
+    distance from the closest already-chosen centroid. This spreads the initial
+    centroids out, which converges faster and more reliably than a uniform draw.
+    Ref: Arthur, D., & Vassilvitskii, S. (2006). "k-means++: the advantages of
+    careful seeding."
+
+    Each step draws `n_local_trials` candidates from that distribution rather than
+    one, and keeps whichever candidate leaves the lowest total inertia. Drawing a
+    single candidate leaves the seeding at the mercy of one unlucky draw: on
+    well-separated blobs it lands in a bad local optimum on roughly a third of
+    seeds, and recovering from that needs about ten restarts. Greedy selection
+    removes those failures at a cost of one extra `(num_samples, n_local_trials)`
+    distance block per step. Note that trials are not a substitute for restarts --
+    every candidate is scored against the same already-chosen prefix, so no number
+    of trials can undo a bad early commitment.
+    """
+    num_samples, num_features = X.shape
+    key, subkey = jr.split(key)
+    first = jr.choice(subkey, X)
+    centroids = jnp.zeros((num_clusters, num_features), X.dtype).at[0].set(first)
+    # Distance from every sample to its closest chosen centroid, carried forward
+    # rather than recomputed. Recomputing costs a (num_samples, num_clusters) block
+    # per step; carrying it costs a (num_samples, n_local_trials) block, which is
+    # smaller as soon as n_local_trials < num_clusters and does not grow with k.
+    closest = _squared_distances(X, first[None, :])[:, 0]
+
+    def step(carry, i):
+        """Draw several candidate centroids and keep the one minimizing inertia."""
+        centroids, closest, key = carry
+        key, subkey = jr.split(key)
+        total = jnp.sum(closest)
+        # If every sample already sits on a centroid, fall back to a uniform draw.
+        probs = jnp.where(
+            total > 0,
+            closest / jnp.where(total > 0, total, 1.0),
+            jnp.full((num_samples,), 1.0 / num_samples, closest.dtype),
+        )
+        candidate_ids = jr.choice(subkey, num_samples, shape=(n_local_trials,), p=probs)
+        candidates = X[candidate_ids]
+        # Column j holds what `closest` would become if candidate j were chosen.
+        distances = jnp.minimum(closest[:, None], _squared_distances(X, candidates))
+        best = jnp.argmin(jnp.sum(distances, axis=0))
+        return (centroids.at[i].set(candidates[best]), distances[:, best], key), None
+
+    (centroids, _, _), _ = lax.scan(step, (centroids, closest, key), jnp.arange(1, num_clusters))
+    return centroids
+
+
+@partial(jit, static_argnames=("k", "max_iters", "n_init", "n_local_trials"))
+def kmeans(
+    X: Float[Array, "num_samples num_features"],
+    k: int,
+    key: PRNGKeyT,
+    max_iters: int = 100,
+    tol: Scalar = 1e-6,
+    n_init: int = 3,
+    n_local_trials: Optional[int] = None,
+) -> KMeansState:
+    """Cluster `X` into `k` groups with Lloyd's algorithm and k-means++ seeding.
+
+    Runs `n_init` independent restarts and returns the one with the lowest inertia,
+    because a single restart can settle in a poor local optimum. Restarts run
+    sequentially via `lax.map` rather than as one vectorized batch: batching
+    materializes a `(n_init, num_samples, k)` distance temporary, whereas
+    `lax.map` stacks only each restart's returned state, which is dominated by
+    the `(num_samples,)` assignments whenever `num_samples` exceeds
+    `k * num_features`. Peak memory therefore grows roughly `k` times more
+    slowly in `n_init`. On CPU this is also faster, since each restart exits at
+    its own convergence instead of the whole batch running until the slowest
+    one converges. Greedy seeding widens that gap rather than closing it: it
+    makes most restarts converge in a handful of iterations while leaving the
+    occasional unlucky one slow, so per-restart iteration counts spread out
+    (measured at N=100k, k=10: [3, 3, 4, 4, 4, 5, 5, 48, 48, 91]) and a batched
+    `while_loop`, which must run every lane until the slowest lane stops,
+    wastes proportionally more work.
+
+    `n_init` defaults to 3 rather than 1 because greedy seeding does not make
+    restarts redundant. All of a step's candidates are scored against the same
+    already-chosen centroids, so no number of trials can undo a bad early
+    commitment; only a fresh restart resamples it. Measured over 20 seeds, one
+    restart still reaches a bad optimum on roughly half of small problems
+    (N=1000, D=2, k=5), while three restarts match or beat the quality of the
+    ten restarts this previously defaulted to, on every workload tested.
+
+    Args:
+        X: samples to cluster.
+        k: number of clusters. Static: changing it triggers recompilation.
+        key: random seed for k-means++ initialization.
+        max_iters: cap on Lloyd iterations per restart. Static.
+        tol: stop once an iteration improves inertia by no more than this.
+        n_init: number of independent restarts. Static.
+        n_local_trials: candidate centroids evaluated per k-means++ step. Defaults to
+            `2 + int(log(k))`. Higher values improve seeding with diminishing returns
+            and are not a substitute for `n_init`. Static.
+
+    Returns:
+        The best `KMeansState` across restarts.
+    """
+
+    trials = 2 + int(math.log(k)) if n_local_trials is None else n_local_trials
+
+    def single_run(key: PRNGKeyT) -> KMeansState:
+        """Run one restart of Lloyd's algorithm from a k-means++ seeding."""
+
+        def cond(carry):
+            """Continue while inertia is still improving by more than tol."""
+            _, previous_inertia, inertia, i = carry
+            return (i < max_iters) & (previous_inertia - inertia > tol)
+
+        def body(carry):
+            """Run one Lloyd iteration: reassign samples, then recompute centroids."""
+            centroids, _, inertia, i = carry
+            assignments = _assign(X, centroids)
+            new_centroids = _update_centroids(X, assignments, k, centroids)
+            return new_centroids, inertia, _inertia(X, new_centroids), i + 1
+
+        initial = _kmeans_plusplus(key, X, k, trials)
+        centroids = _update_centroids(X, _assign(X, initial), k, initial)
+        carry = (centroids, jnp.inf, _inertia(X, centroids), 1)
+        centroids, _, inertia, n_iter = lax.while_loop(cond, body, carry)
+        return KMeansState(centroids, _assign(X, centroids), inertia, n_iter)
+
+    restarts = lax.map(single_run, jr.split(key, n_init))
+    best = jnp.argmin(restarts.inertia)
+    return KMeansState(
+        restarts.centroids[best],
+        restarts.assignments[best],
+        restarts.inertia[best],
+        restarts.n_iter[best],
+    )
diff --git a/dynamax/utils/cluster_test.py b/dynamax/utils/cluster_test.py
new file mode 100644
index 00000000..8e3d5025
--- /dev/null
+++ b/dynamax/utils/cluster_test.py
@@ -0,0 +1,151 @@
+"""Tests for the k-means clustering utilities."""
+
+import jax.numpy as jnp
+import jax.random as jr
+from jax import jit, vmap
+
+from dynamax.utils.cluster import kmeans
+
+
+def test_kmeans_recovers_toy_optimum():
+    """Recovers the known optimal 2-means solution on a 4-point toy example."""
+    x = jnp.array([[0.0, 0.0], [0.5, 0.0], [0.5, 1.0], [1.0, 1.0]])
+    state = kmeans(x, 2, jr.PRNGKey(101))
+    # Compare permutation-invariantly by sorting centroids on their y coordinate.
+    centroids = state.centroids[jnp.argsort(state.centroids[:, 1])]
+    assert jnp.allclose(centroids, jnp.array([[0.25, 0.0], [0.75, 1.0]]), atol=1e-5)
+    assert jnp.allclose(state.inertia, 0.25, atol=1e-5)
+
+
+def test_kmeans_recovers_separated_clusters():
+    """Assignments perfectly partition well-separated blobs."""
+    key = jr.PRNGKey(0)
+    means = jnp.array([[-6.0, -6.0], [0.0, 0.0], [6.0, 6.0]])
+    x = jnp.concatenate([m + 0.3 * jr.normal(k, (40, 2)) for m, k in zip(means, jr.split(key, 3))])
+    state = kmeans(x, 3, jr.PRNGKey(1))
+    true_labels = jnp.repeat(jnp.arange(3), 40)
+    # Each true blob must map onto exactly one predicted label.
+    for blob in range(3):
+        assert jnp.unique(state.assignments[true_labels == blob]).size == 1
+    assert jnp.unique(state.assignments).size == 3
+
+
+def test_kmeans_empty_clusters_do_not_produce_nans():
+    """k exceeding the number of distinct points must not yield NaN centroids."""
+    x = jnp.ones((4, 2))
+    state = kmeans(x, 3, jr.PRNGKey(0))
+    assert not jnp.any(jnp.isnan(state.centroids))
+    assert jnp.allclose(state.inertia, 0.0)
+    assert int(state.n_iter) < 10  # terminates promptly instead of spinning to max_iters
+
+
+def test_kmeans_one_cluster_per_sample():
+    """With k equal to the sample count, every sample becomes its own centroid."""
+    x = jnp.array([[0.0, 0.0], [0.1, 0.0], [10.0, 10.0]])
+    state = kmeans(x, 3, jr.PRNGKey(0))
+    assert jnp.unique(state.assignments).size == 3
+    assert jnp.allclose(state.centroids[state.assignments], x, atol=1e-5)
+    assert jnp.allclose(state.inertia, 0.0, atol=1e-5)
+
+
+def test_kmeans_one_dimensional_data():
+    """Works on (N, 1) data, the shape GammaHMM passes in."""
+    x = jnp.concatenate([jr.normal(jr.PRNGKey(1), (50, 1)) + 5, jr.normal(jr.PRNGKey(2), (50, 1)) - 5])
+    state = kmeans(x, 2, jr.PRNGKey(0))
+    centroids = jnp.sort(state.centroids.ravel())
+    assert centroids[0] < 0 < centroids[1]
+    assert jnp.unique(state.assignments).size == 2
+
+
+def test_kmeans_inertia_matches_assignments():
+    """Reported inertia equals the sum of squared distances to assigned centroids."""
+    x = jr.normal(jr.PRNGKey(4), (100, 3))
+    state = kmeans(x, 4, jr.PRNGKey(5))
+    recomputed = jnp.sum((x - state.centroids[state.assignments]) ** 2)
+    assert jnp.allclose(state.inertia, recomputed, rtol=1e-5)
+
+
+def test_kmeans_is_deterministic_given_key():
+    """The same key produces identical results."""
+    x = jr.normal(jr.PRNGKey(6), (80, 2))
+    a = kmeans(x, 3, jr.PRNGKey(7))
+    b = kmeans(x, 3, jr.PRNGKey(7))
+    assert jnp.allclose(a.centroids, b.centroids)
+    assert jnp.all(a.assignments == b.assignments)
+
+
+def test_kmeans_restarts_beat_single_init():
+    """Restarts still buy quality that better seeding alone cannot.
+
+    Greedy k-means++ scores every candidate against the centroids already chosen,
+    so it cannot recover from a bad early commitment; only a fresh restart
+    resamples one. On this fixture a single initialization reaches a bad optimum
+    on most seeds regardless of how many candidates each step considers, which is
+    why n_init stays above 1 by default.
+    """
+    means = 4.0 * jr.normal(jr.PRNGKey(11), (5, 2))
+    idx = jr.randint(jr.PRNGKey(12), (1000,), 0, 5)
+    x = means[idx] + 0.7 * jr.normal(jr.PRNGKey(13), (1000, 2))
+    single = min(float(kmeans(x, 5, jr.PRNGKey(s), n_init=1).inertia) for s in range(6))
+    many = min(float(kmeans(x, 5, jr.PRNGKey(s), n_init=10).inertia) for s in range(6))
+    assert many <= single
+    assert many < 900.0  # the good optimum is 883.9; the bad one is ~1048
+
+
+def test_kmeans_single_init_finds_optimum_on_separated_blobs():
+    """A single initialization must reach the good optimum on well-separated blobs.
+
+    Drawing one candidate per k-means++ step lands in a bad local optimum on seeds
+    0, 5 and 7 of this fixture (inertia ~1044, ~1229, ~1044 against an optimum of
+    184.3), which is what forced a high n_init default. Drawing several candidates
+    per step and keeping the one that most reduces inertia fixes all of them.
+    """
+    key = jr.PRNGKey(0)
+    means = jnp.array([[-4.0, -4.0], [0.0, 0.0], [4.0, 4.0], [8.0, -4.0]])
+    x = jnp.concatenate([m + 0.6 * jr.normal(k, (60, 2)) for m, k in zip(means, jr.split(key, 4))])
+    for seed in range(8):
+        inertia = kmeans(x, 4, jr.PRNGKey(seed), n_init=1).inertia
+        assert inertia < 200.0, f"seed {seed} landed in a bad local optimum (inertia {inertia:.1f})"
+
+
+def test_kmeans_n_local_trials_is_configurable():
+    """n_local_trials is honored, and one trial reproduces plain k-means++."""
+    key = jr.PRNGKey(0)
+    means = jnp.array([[-4.0, -4.0], [0.0, 0.0], [4.0, 4.0], [8.0, -4.0]])
+    x = jnp.concatenate([m + 0.6 * jr.normal(k, (60, 2)) for m, k in zip(means, jr.split(key, 4))])
+    # Seed 0 is one of the seeds a single candidate per step gets wrong.
+    assert kmeans(x, 4, jr.PRNGKey(0), n_init=1, n_local_trials=1).inertia > 200.0
+    assert kmeans(x, 4, jr.PRNGKey(0), n_init=1, n_local_trials=16).inertia < 200.0
+
+
+def test_kmeans_is_jittable_and_vmappable():
+    """Works under jit and vmap over a batch of datasets."""
+    x = jr.normal(jr.PRNGKey(8), (4, 60, 2)) * 2
+    keys = jr.split(jr.PRNGKey(9), 4)
+    batched = vmap(lambda d, k: kmeans(d, 3, k))(x, keys)
+    assert batched.centroids.shape == (4, 3, 2)
+    assert batched.assignments.shape == (4, 60)
+    assert jnp.all(jnp.isfinite(batched.inertia))
+
+    single = jit(lambda d, k: kmeans(d, 3, k))(x[0], keys[0])
+    assert jnp.allclose(single.centroids, batched.centroids[0])
+
+
+def test_kmeans_stable_at_large_offset():
+    """Squared distances stay non-negative and assignments stay correct far from the origin.
+
+    Regression test: expanding ||x - c||^2 to ||x||^2 - 2 x.c + ||c||^2 without first
+    centering on the data is not offset-stable in float32 and can return negative
+    "squared" distances, which silently corrupts both clustering and n_init restart
+    selection (kmeans picks the restart with the lowest, possibly negative, inertia).
+    """
+    key = jr.PRNGKey(0)
+    means = jnp.array([[-6.0, -6.0], [0.0, 0.0], [6.0, 6.0]])
+    x = jnp.concatenate([m + 0.3 * jr.normal(k, (40, 2)) for m, k in zip(means, jr.split(key, 3))])
+    x_shifted = x + 1e5
+    state = kmeans(x_shifted, 3, jr.PRNGKey(1))
+    assert state.inertia >= 0
+    true_labels = jnp.repeat(jnp.arange(3), 40)
+    for blob in range(3):
+        assert jnp.unique(state.assignments[true_labels == blob]).size == 1
+    assert jnp.unique(state.assignments).size == 3
diff --git a/pyproject.toml b/pyproject.toml
index 31891f38..1693c73a 100644
--- a/pyproject.toml
+++ b/pyproject.toml
@@ -12,7 +12,6 @@ dependencies = [
     "tfp-nightly",
     "fastprogress",
     "optax",
-    "scikit-learn",
     "jaxtyping",
     "typing-extensions",
     "numpy"
@@ -46,7 +45,8 @@ notebooks = [
     "flax",
     "blackjax",
     "graphviz",
-    "scipy"
+    "scipy",
+    "scikit-learn"
 ]
 
 doc = [
@@ -56,6 +56,7 @@ doc = [
     "blackjax",
     "graphviz",
     "scipy",
+    "scikit-learn",
     "sphinx",
     "sphinx-autobuild",
     "sphinx_autodoc_typehints",
@@ -81,6 +82,7 @@ dev = [
     "blackjax",
     "graphviz",
     "scipy",
+    "scikit-learn",
     "sphinx",
     "sphinx-autobuild",
     "sphinx_autodoc_typehints",