mirror of
https://github.com/ROCm/jax.git
synced 2025-04-19 05:16:06 +00:00
ff18dedf99
Once the compiler supports it, we can make it public and allow users to pass those values. Right now, only `major_to_minor` is supported. But a valid question is why even keep them as arguments in the constructor? It's because we need to translate `PjRtLayout` which we get from the executable to `DeviceLocalLayout` and preserve the `tiling` and `sub_byte_element_size_in_bits` info that we get from the compiler. This has helped catch bugs before when the compiler was not doing the right thing in layout propagation pass. PiperOrigin-RevId: 651644934