bottleneck-transformer-flax

Depends On

• jax (jax)
• jaxlib (jaxlib)
• flax (flax)
• einops (einops>=0.3)