jax

0cd5ca77 - [Mosaic GPU] Lock down 8-bit pointwise ops (except bitwise and cmp) + add lshift/rshift

[Mosaic GPU] Lock down 8-bit pointwise ops (except bitwise and cmp) + add lshift/rshift 8-bit ops generally have poor hardware support and sometimes produce really cryptic LLVM or MLIR errors (e.g. when trying to add 8-bit floats). PiperOrigin-RevId: 857122323