用JAX快速实现RWKV LLM的进化算法

jax.vmap 和 jax.random.key 是我能够在 1 天内为 RWKV LLM 实现 EGGROLL 的原因。我 literally 某天醒来（8月13日）想“等等，我能不能直接对 RWKV 做进化”，然后就成功了。