hyena-jax

Port of Hyena Hierarchy to JAX

Original: https://github.com/HazyResearch/safari/blob/main/standalone_hyena.py

Flax implementation : https://github.com/Ryu1845/hyena-jax/blob/main/hyena_flax.py