jax

46e65b59 - [pallas] Added API docs for Triton and Mosaic GPU backends

[pallas] Added API docs for Triton and Mosaic GPU backends I've left the TPU backend docs a stub for now. Hopefully, someone working on Pallas TPU can fill them in later.