Efficient Mixture of Experts implementation?

[SamKG]

Hello,

Is there an efficient implementation of Mixture of Experts models (e.g. Mixtral) anywhere that uses flax/jax?
I can't seem to find one, though I'd be happy to try writing one myself if there's interest.

As an aside: is block-sparse matmul efficient in Jax/Flax?

[cgarciae]

Checkout the moe implementation from Flaxformers.