Pallas Puzzles

tl;dr. vorushin/pallas_puzzles - JAX’s kernel language puzzles Pallas is JAX’s kernel language for writing custom operations that run on TPU. Pallas for TPUs is what Triton1 is for GPUs. The kernels are written to speed up model training and inference - the kernel languages provide low-level access to the hardware, allowing you to perform optimizations outside of the compiler’s reach. Tri Dao2…