jax

acd0e497 - [pallas:mosaic_gpu] `GPUBlockSpec` no longer accepts `swizzle`

[pallas:mosaic_gpu] `GPUBlockSpec` no longer accepts `swizzle` It was previously possible to pass `swizzle` both directly and via `transforms`. This change eliminates the ambiguity at a slight downgrade to ergonomics. PiperOrigin-RevId: 684797980