faiss
jax[cpu]
flax