All modules for which code is available

• jax.numpy
• jaxgarden.attention.multi_head_attention
• jaxgarden.functional.attention
• jaxgarden.models.base
• jaxgarden.models.gemma2
• jaxgarden.models.gemma3
• jaxgarden.models.generation_utils
• jaxgarden.models.llama
• jaxgarden.models.modernbert