This repository was archived by the owner on Aug 25, 2026. It is now read-only.

Repository navigation

All pull requests

Labels (9)Milestones (0)

Add Gemma-4 support
#139JEM-Mosig opened Aug 10, 2026Contributor
feat: Add Geometric Sparse Attention (AETHER)
#134teerthsharma opened Jan 20, 2026
1
Add py.typed marker.
#130Jammf opened Oct 17, 2025
Fix dtype type annotations to use jax.typing.DTypeLike.
#129Jammf opened Oct 16, 2025
Use jax.tree_util.treedef_is_leaf for readability
#128findmyway opened Aug 19, 2025
Feature Request: Add convolution layers
#120AntoinePlumerault opened May 31, 2025
41
Experiment with using out_axes=-1 inside nmap's vmap
#99shoyer opened Nov 25, 2024