rocm_jax

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

History

Peter Hawkins 3fa557289a Port tests away from setUpClass and setUpModule to setUp alone.

This change prepares for upcoming changes in which we run tests in parallel using threads, which we are doing partially to test free threading but also partially to speed up TPU tests via thread-parallelism.

If independent tests run in parallel in no particular order, there's no natural scope around which to call setUpClass or SetUpModule. But for JAX tests this never seems necessary: we can just do the same work in setUp() or do it globally.

PiperOrigin-RevId: 713296722

2025-01-08 08:14:50 -08:00

BUILD

[Mosaic GPU] Split layout inference and dialect lowering files and tests.

2024-12-11 07:31:34 -08:00

flash_attention_test.py

[Mosaic GPU] Only run tests requiring sm90a on Hopper

2024-11-11 12:02:48 -05:00

gpu_dialect_test.py

[Mosaic GPU] Add a lowering for simple async_load and async_store ops.

2024-12-20 09:38:13 -08:00

gpu_layout_inference_test.py

[Mosaic GPU] Commit to using Vectors everywhere (and no Tensors).