[Mosaic TPU] Add lowering rule for `lax.approx_max_k` in Pallas Mosaic. First reduce the reduction dimension to some number of bins, then do top_k.
PiperOrigin-RevId: 972170791
Y
Yue Sheng committed
2d2fb43ef29244ab0486db416ca5de83f177cccd
Parent: bbbfdcd
Committed by jax authors <google-ml-automation@google.com>
on 8/27/2026, 9:52:53 PM