jax

d751a3f7 - [Mosaic GPU] Fix layout inference for `tcgen05_mma` to be compatible with user-specified tile sizes.

[Mosaic GPU] Fix layout inference for `tcgen05_mma` to be compatible with user-specified tile sizes. Previously, we chose the largest possible swizzle value in order to determine the tile size along the minormost dimension of the `tcgen05_mma`'s SMEM operands. While using the largest tiling possible is a reasonable thing to do, it is not strictly necessary for correctness. As a result, it is possible for a user to specify a smaller tiling, and still produce a valid program. In this case, our strict assignment logic immediately made the system unsatisfiable, despite the existence of a valid solution that would satisfy the user's constraints. This change replaces the assignment with a new constraint `IsValidMmaTiling`. This allows applying the user-specified constraint if it exists, and also allows us to still conjure the largest possible tiling if none is specified. PiperOrigin-RevId: 863444916