> 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