jax

94748a15 - [Mosaic GPU] Fix bug in the implementation of `tile_strides` in the presence of duplicate strides.

[Mosaic GPU] Fix bug in the implementation of `tile_strides` in the presence of duplicate strides. Previously, we would reorder tile sizes by looking up the indices of their corresponding strides. This did not work correctly in the case of duplicate strides, since the first index of the stride within the list of strides would be returned. We now `zip` the strides with the tiling, and use a composite key `(stride, tiling)` to perform the lookups, resolving the ambiguity. PiperOrigin-RevId: 899615083