Commits ROCm/jax - SemanticDiff

Rollback due to downstream test failures

Google-ML-Automation committed 143 days ago

8d2c8ad5

Add `AbstractDevice` to `AbstractMesh` to capture specific device properties; `device_kind` and `num_tpu_cores`.

yashk2810 committed 143 days ago

8be91cde

Merge pull request #31131 from mattjj:discharge-state-caching

Google-ML-Automation committed 143 days ago

5977fc2e

[mutable-arrays] add discharge_state2 for ClosedJaxprs, with caching

mattjj committed 143 days ago

5ff55080

Expose a standard function to insert pvary in `jax.random` APIs that we can reuse.

yashk2810 committed 143 days ago

9d1661ed

[Mosaic GPU] Use a custom pass to compile PTX to CUBIN.

bchetioui committed 143 days ago

87d5de4f

Make jax.lax.empty support uninitialized values on TPU and expose the API as `jax.lax.empty`

yashk2810 committed 143 days ago

99d4cc89

[Mosaic GPU] Simplify tests by removing intermediate copies to SMEM.

allanrenucci committed 143 days ago

b224f338

Produces better error message on tree mismatch, before we would dump potentially huge subtrees into the log which would often get truncated and would be missing needed information (which keys actually mismatch), now we only include the keys.

marksandler2 committed 143 days ago

f7f021e5

[Mosaic GPU] Stop exporting old-style custom calls.

apaszke committed 143 days ago

591dda01

[Mosaic GPU][NFC] Simplify usages of `DialectBarrierRef` in MGPU dialect tests.

allanrenucci committed 143 days ago

602fc1e8

[Pallas:MGPU] Increase test shard count to avoid ASAN/TSAN timeouts

apaszke committed 143 days ago

6e91e96e

[Pallas/interpreter] Raise exception when attempting to (de)reference memory spaces `HBM` or `ANY` from inside a kernel.

Google-ML-Automation committed 143 days ago

9c4f04fe

[Pallas:MGPU] Expose gathers in plgpu.copy_gmem_to_smem

apaszke committed 143 days ago

ab57051f

Add env flag to fail when testing collective matmul without nvshmem support.

Google-ML-Automation committed 143 days ago

417da5d0

[Mosaic GPU] Add support for transforms in gather TMA

apaszke committed 143 days ago

180c448d

[Mosaic GPU] Add a layout inference rule for `vector.ShapeCastOp`.

bchetioui committed 143 days ago

ab57bcb7

Fix usage of `tcgen05.commit_arrive` in collective MMA test.

allanrenucci committed 143 days ago

d98d979a

[Mosaic GPU] Simplify inline_asm for optimization_barrier

apaszke committed 143 days ago

f0352e1f

[Mosaic GPU] Respect `packing` for TMEM scratch allocation.

allanrenucci committed 143 days ago

bba6c7e3

Update XLA dependency to use revision http://github.com/openxla/xla/commit/53db8f76024aeba52d5b2b06295adc793801d10f

Google-ML-Automation committed 143 days ago

b3d41ea4

Update XLA dependency to use revision http://github.com/openxla/xla/commit/28f1592fd901264e429bebfa8b7e0ecd0ec22112

Google-ML-Automation committed 144 days ago

16132622

Shape validation in _array_from_process_local_data` for single-process case.

marksandler2 committed 145 days ago

3688db99

Update XLA dependency to use revision http://github.com/openxla/xla/commit/92f7b5952dd585c5be17c9a5caad27407005b513

Google-ML-Automation committed 145 days ago

c8091135

Fix tuple comparison in jax.random vma comparison.

pschuh committed 146 days ago

66f06c85

Merge pull request #31088 from jakevdp:mypy-pre-commit

Google-ML-Automation committed 146 days ago

ffeddf7b

[pre-commit] bump mypy to v1.17.1

jakevdp committed 146 days ago

6f72a253

Reverts 28140bb7cd7c247ebb9ab8d29a34d29d1ec41fb2

yashk2810 committed 146 days ago

920ae39f

Add `core.pvary` at the right places in `jax.random.gamma` implementation so that we don't raise internal errors. Fixes https://github.com/jax-ml/jax/issues/31080

pschuh committed 146 days ago

6ed34cd9

Merge pull request #31087 from jakevdp:pre-commit-version

Google-ML-Automation committed 146 days ago

8f02f7f9