Commits ROCm/jax - SemanticDiff

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

bchetioui committed 172 days ago

87d5de4f

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

yashk2810 committed 172 days ago

99d4cc89

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

allanrenucci committed 172 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 172 days ago

f7f021e5

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

apaszke committed 172 days ago

591dda01

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

allanrenucci committed 172 days ago

602fc1e8

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

apaszke committed 172 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 172 days ago

9c4f04fe

[Pallas:MGPU] Expose gathers in plgpu.copy_gmem_to_smem

apaszke committed 172 days ago

ab57051f

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

Google-ML-Automation committed 172 days ago

417da5d0

[Mosaic GPU] Add support for transforms in gather TMA

apaszke committed 172 days ago

180c448d

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

bchetioui committed 172 days ago

ab57bcb7

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

allanrenucci committed 172 days ago

d98d979a

[Mosaic GPU] Simplify inline_asm for optimization_barrier

apaszke committed 172 days ago

f0352e1f

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

allanrenucci committed 172 days ago

bba6c7e3

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

Google-ML-Automation committed 172 days ago

b3d41ea4

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

Google-ML-Automation committed 173 days ago

16132622

Shape validation in _array_from_process_local_data` for single-process case.

marksandler2 committed 173 days ago

3688db99

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

Google-ML-Automation committed 174 days ago

c8091135

Fix tuple comparison in jax.random vma comparison.

pschuh committed 174 days ago

66f06c85

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

Google-ML-Automation committed 174 days ago

ffeddf7b

[pre-commit] bump mypy to v1.17.1

jakevdp committed 174 days ago

6f72a253

Reverts 28140bb7cd7c247ebb9ab8d29a34d29d1ec41fb2

yashk2810 committed 174 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 174 days ago

6ed34cd9

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

Google-ML-Automation committed 174 days ago

8f02f7f9

Make sure `use_abstract_mesh` can only override `axis_types` i.e. it **cannot** override `axis_names` and `axis_sizes`. Also delete `use_auto_axes` and `use_explicit_axes` API. Users can just use `use_abstract_mesh` to get the same functionality.

yashk2810 committed 174 days ago

28140bb7

[pre-commit] bump pre-commit-hooks version

jakevdp committed 174 days ago

a1db7ba7

Merge pull request #31086 from jakevdp:update-ruff

Google-ML-Automation committed 174 days ago

3ed6b4c0

[pre-commit] update ruff to v0.12.9

jakevdp committed 174 days ago

1074ba57

Fix typos

ZacCranko committed 174 days ago

a2127d0c