einops
jaxtyping
schedulefree
liger-kernel