Davis Yoshida

lorax. LoRA for arbitrary JAX models and functions

qax. If it quacks like a tensor...

tf2-gradient-checkpointing. Simple gradient checkpointing for eager mode execution

haiku-mup. A port of muP to JAX/Haiku

abnormal-floats. Code for the note "NF4 Isn't Information Theoretically Optimal (and that's Good)

easy-lora-and-gptq. JAX notebook showing how to LoRA + GPTQ arbitrary models

jax-gptq. JAX implementation of GPTQ quantization algorithm

jax-dropless-moe. WIP implementation of block-sparse dropless MoE in JAX

gpt-2-haiku. My port of GPT-2 to JAX/haiku. You probably want the HuggingFace FLAX one instead.

360blockscope. Block scoping in python (ironically)

asyncbots. A framework for simplifying writing RTM bots for Slack.