Palladium
I'm currently working on a JAX Pallas demo for Apple GPU that works by translating jaxprs into Metal shading language, compiling that, and then executing with buffers from Python exchanged with nanobind. I planned this as just an exercise for me to learn more about MSL, GPU kernel programming, and to shake off the rust on numerics accumulated since studies, but if anyone wants to try it, you can find it at https://github.com/nicholasjng/palladium.