jax
f9a10ff0 - [Mosaic GPU] Fix bugs in `transform_type` implementation.

Commit
14 days ago
[Mosaic GPU] Fix bugs in `transform_type` implementation. The permutation logic was inverted, and we were missing a divisor in the computation of strides. This was not caught earlier, because other tests focus on tiling up to 2 dimensions. A map permuting 2 dimensions is always its own inverse---making the permutation logic coincidentally correct. Similarly, since the loop computing the strides between the indices of the tiles runs `len(tiling) - 1 = 1` times, and the missing divisor is `1` in this case, the result is also coincidentally correct. We now add property-based tests in order to ensure the correctness of the implementation. PiperOrigin-RevId: 868620175
Author
Parents
Loading