rocm_jax

mirror of https://github.com/ROCm/jax.git synced 2025-04-19 05:16:06 +00:00

History

George Necula 9b35b760ce [pallas] Enable check for GPU lowering that tensor sizes are power of 2

Triton has a restriction that all operations have arguments and results
that are tensor whose size is a power of 2. Added a lowering check
for this. Without this, when we violate the condition we get an
unfriendly crash.

PiperOrigin-RevId: 659483450

2024-08-05 02:34:21 -07:00

tpu

[pallas] Add error message for block_shapes of rank less than 1.

2024-08-01 09:15:07 -07:00

CHANGELOG.md

[pallas] Add error message for block_shapes of rank less than 1.

2024-08-01 09:15:07 -07:00

design.md

Changed `pl.BlockSpec to accept block_shape before index_map`

2024-07-01 14:26:08 -07:00

grid_blockspec.md

[pallas] Enable check for GPU lowering that tensor sizes are power of 2