About
Felafax builds RL environments that make models better at writing JAX code and Pallas kernels.
We were early believers in JAX and TPUs. We're ex‑Google and Meta engineers who did a lot of work on TPUs before they were popular (Hacker News [1] & [2], AMD blog). Now we're getting back to them.
Our environments start from work we used to do by hand: porting models from PyTorch to JAX, and making them fast on TPU.
Want a sample environment? Email founders@felafax.ai
Environments
PyTorch → JAX
The agent ports a popular open-source model from PyTorch to JAX. The grader checks for correctness and clean JAX code. We've trained small LLMs with it ourselves and seen them improve on JAXBench.
Request a sample →Pallas kernels
The agent writes optimized Pallas kernels, such as grouped matmul and sliding window attention.
Request a sample →TPU megakernels
Coming soon.
Prior work
Team
co-founders (and twin brothers!)
nithin
Previously ML engineer at Google/YouTube and NVIDIA.
nikhil
Previously engineer at Meta and Microsoft.