Comment by yarri

I assume the Gemini results are JAX/PAX-ML/Pallas improvements for TPUs so would look there for recent PRs