faiss-cpu
jax[cpu]
flax