ROCm/jax

Pull Requests Commits

Enabled ROCm devices to default to Jacobi SVD on smaller matrices.

tsrw2048 committed 2 days ago

29d5a28c

Merge pull request #34488 from jakevdp:mode-axis

Google-ML-Automation committed 2 days ago

b8da25b9

Merge pull request #34485 from jakevdp:metal-docs

Google-ML-Automation committed 2 days ago

2edcae89

Try to fix flaky `test_unknown_platform_error` test.

mwhittaker committed 2 days ago

b64d4e23

Remove mentions of jax-metal

jakevdp committed 2 days ago

fc5fccff

jax.scipy.stats.mode: accept negative indices

jakevdp committed 2 days ago

28843f7e

Automated Code Change

Google-ML-Automation committed 2 days ago

f49e5d03

[Pallas/Mosaic GPU] Annotate test that modifies environment with `jtu.thread_unsafe_test`.

bchetioui committed 3 days ago

ce373b68

[Mosaic GPU][NFC] Simplify allocation in `MosaicGpuCompile`.

allanrenucci committed 3 days ago

bf0555ac

[XLA:MGPU] Enforce shape.size() < tiled_tiling_shape.size().

golechwierowicz committed 3 days ago

d6e76afa

[XLA:MGPU] Port remove_dimension and reduce to C++.

golechwierowicz committed 3 days ago

36e3dd61

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

Google-ML-Automation committed 3 days ago

a5ca98f4

[Mosaic GPU] Statically check for SMEM out of bound slicing.

allanrenucci committed 3 days ago

7de23ec8

Use maxsize=None with trace_to_jaxpr's weakref_lru_cache to get more cache hits. This also restores previous behavior of lu.cache.

yashk2810 committed 3 days ago

d3c26091

[Pallas:MGPU] Add support for cross-warp reductions under WG semantics.

allanrenucci committed 3 days ago

ba9621cf

Merge pull request #34288 from olupton:smem-limit

Google-ML-Automation committed 3 days ago

1033308a

Artificially limit SMEM capacity for tests

olupton committed 3 days ago

490b6472

[XLA:MGPU] Port registers_shape, shape_from_registers_shape, registers_element_type methods.

golechwierowicz committed 3 days ago

ee387893

[Pallas/Mosaic GPU] Expose `plgpu.wgmma_accumulator_load`.

bchetioui committed 3 days ago

08a8bf71

[Mosaic GPU] Allow load/store_tiled with refs where the tiling rank is smaller than the logical rank

apaszke committed 3 days ago

cb24d7fc

[XLA:MGPU] Port warp_indices and lane_indices to C++.

golechwierowicz committed 4 days ago

0cf96f75

Merge pull request #34402 from gnecula:clean_dynamic_shape_replacement

Google-ML-Automation committed 4 days ago

64301806

Merge pull request #34425 from andportnoy:aportnoy/mosaic-gpu-fence-proxy-async

Google-ML-Automation committed 4 days ago

c32e170b

[pallas] Minor cleanup of typing annotations for dynamic_shape_replacement_fn

gnecula committed 4 days ago

a51d661b

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

Google-ML-Automation committed 4 days ago

7c5c109d

Update XLA dependency to use revision http://github.com/openxla/xla/commit/7d9b40ba458f389d978603fb78d7fe4bff02100c

Google-ML-Automation committed 5 days ago

ff99faf2

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

Google-ML-Automation committed 6 days ago

c6568036

Remove ref cycles introduced by self-calling nested functions.

ICGog committed 6 days ago

15199269

[jax] bound traceback of pallas_call to within the kernel.

cezheng committed 6 days ago

fe359a15

Clean up a few TODOs in pjit after FlatTree change

yashk2810 committed 6 days ago

aeccdff4

Older