About
Felafax builds RL environments that make models better at writing JAX code and Pallas kernels.
We were early believers in 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
nithin
Previously ML engineer at Google/YouTube and NVIDIA.
nikhil
Previously engineer at Meta and Microsoft.