9 Commits

Author SHA1 Message Date
Charles Hofer
c7b407c9f0 Merge branch 'rocm-main' into ci-upstream-sync-151_1 2025-03-18 15:27:35 +00:00
Sergei Lebedev
0ff234049b Removed trivial docstrings from JAX tests
These docstrings do not make the tests any more clear and typically just duplicate the test module name.

PiperOrigin-RevId: 737611977
2025-03-17 07:49:37 -07:00
Charles Hofer
47580efda5 Merge branch 'rocm-main' into ci-upstream-sync-98_1 2025-01-28 21:18:47 +00:00
Charles Hofer
41ab12bf8d Skip PallasCallRemoteDMAInterpretTest.test_interpret_remote_dma_ppermute for failing on ROCm 2025-01-27 20:46:14 +00:00
Jake VanderPlas
8e97272332 tpu_pallas_distributed_test: skip unless on TPU 2025-01-27 12:41:28 -08:00
Adam Paszke
d5ead570bb [Mosaic TPU] Add support for modeling loads/stores and fix minor issues in model extraction
PiperOrigin-RevId: 703102072
2024-12-05 07:06:19 -08:00
Bart Chrzaszcz
801fe87da6 Do not allow None axis names in meshes.
PiperOrigin-RevId: 686557025
2024-10-16 10:32:25 -07:00
Justin Fu
3c7c9ffbbc [Pallas] Correctly handle asymmetrical remote DMA dst_ref indexing in interpret mode.
PiperOrigin-RevId: 658029297
2024-07-31 09:10:53 -07:00
Justin Fu
acacbe297f [Pallas] Move distributed TPU tests to their own file.
PiperOrigin-RevId: 658014149
2024-07-31 08:19:19 -07:00