[emu]

> Such capability can’t even be contemplated on the TPU side because you basically get a black box.

I'll just leave this here: https://jax.readthedocs.io/en/latest/pallas/index.html

[ein0p]

Pallas is very new. Given how difficult these things are to debug and how half assed the XLA tooling generally is, I’d give it at least another year, if not two, before I touch it for anything prod related.