Tokamax: A GPU and TPU Custom Kernel Library

Tokamax is an OpenXLA library providing high-performance custom accelerator kernels for NVIDIA GPUs and Google TPUs. It offers state-of-the-art implementations built on top of JAX and Pallas, along with tooling for users to build and autotune their own custom kernels. As of version 0.0.12, it is still under heavy development, and users should anticipate API changes.