[Mosaic TPU] Add support for bf16 top_k.
PiperOrigin-RevId: 970013890
Y
Yue Sheng committed
2380871fc9fab5c561ba75e0326c868738da9e75
Parent: d12cf04
Committed by jax authors <google-ml-automation@google.com>
on 8/24/2026, 7:40:16 PM