ROCm/jax

Pull Requests Commits

Disable GPU memory preallocation in ROCm pytest

gulsumgudukbay committed 107 days ago

d2d242ad

Merge pull request #35939 from jakevdp:del-arraylike

Google-ML-Automation committed 108 days ago

348bff56

[cleanup] remove unnecessary call to lax.asarray

jakevdp committed 108 days ago

53ee1d0e

Make sure to use `to_ct_spec()` when unsharding Zeros in shard_map_transpose. Also use `to_ct_aval` in _flatten_bwd in custom_vjp

yashk2810 committed 108 days ago

02007235

[Mosaic GPU] Add async_copy_sparse_metadata_smem_to_tmem dialect lowering.

golechwierowicz committed 108 days ago

12a24495

[Mosaic GPU] Test tcgen05_mma_sparse_metadata layout inference when LHS is in TMEM.

golechwierowicz committed 108 days ago

555baca7

[Mosaic GPU][NFC] Cleanup error reporting.

allanrenucci committed 109 days ago

3af20779

[Mosaic GPU] Remove a_sparse_metadata availability check.

golechwierowicz committed 109 days ago

196f85b9

Fix a flake in LoggingTest.test_debug_logging.

hawkinsp committed 109 days ago

3c0df6e3

Disable hypothesis too_slow checks under ASAN/TSAN/MSAN.

hawkinsp committed 109 days ago

054eaacf

[Mosaic GPU] Add Pallas lowering for AsyncStoreSmemToTmem op.

golechwierowicz committed 109 days ago

ed8e9d55

[Mosaic GPU] Make tests involving global state order-independent.

bchetioui committed 109 days ago

a86c180e

[NFC] Reduce nesting with early return in `_core_map_lowering_rule`.

allanrenucci committed 109 days ago

7f265b45

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

Google-ML-Automation committed 109 days ago

e8253978

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

Google-ML-Automation committed 110 days ago

8304fec2

Merge pull request #34332 from jakevdp:searchsorted-hijax

Google-ML-Automation committed 110 days ago

13dcc922

Reduce the number of pytest workers for Linux aarch64 CPU builds.

hawkinsp committed 110 days ago

6463906c

Use a np.printoptions context manager rather than np.set_printoptions in test utility.

hawkinsp committed 110 days ago

5c4dd597

Reverts 3b4bee803de1a83b98cfe1aced47ab58f9bbc0c3

hawkinsp committed 110 days ago

7a68c4b1

jakevdp committed 111 days ago

5d7b5186

validate out_dtype

jakevdp committed 111 days ago

250ced1d

more fixes for sharded computation

jakevdp committed 111 days ago

aad9769e

Use scan rather than fori_loop

jakevdp committed 111 days ago

6ef7e8d8

[hijax] add hijax searchsorted primitive

jakevdp committed 111 days ago

2185cb1c

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

Google-ML-Automation committed 111 days ago

4ea536c2

Merge pull request #35696 from davisyoshida:recip-docstring

Google-ML-Automation committed 111 days ago

ea4b2e7e

davisyoshida committed 111 days ago

a767ad38

Fix reciprocal docstring

davisyoshida committed 111 days ago

f5676668

Automated Code Change

Google-ML-Automation committed 111 days ago

2d5a40fe

Make _get_owner skip sharding_constraints so that the xla_metadata is applied to the op instead of the sharding constraints

yashk2810 committed 111 days ago

a364e282

Older