Trim heavy dependencies: reconsider hard scikit-learn dependency for KMeans-only usage

[kylejcaron]

This appears to be a hard dependency purely for KMeans used in initialization for HMMs. Pulling in sklearn as a top-level dependency brings in a large, non-JAX dependency chain for what's a single clustering call used at init time.

Can an HMM extra be added, and sklearn can be moved over to it?

[slinderman]

Agreed. We wrote a simple kmeans alg in JAX to avoid this dependency, but I don’t think we ever tested it enough to deploy. Finishing that would be cleaner than having another set of requirements for HMMs in my opinion. Should be pretty straightforward.