ROCm/jax

Pull Requests Commits

Update XLA dependency to use revision

Google-ML-Automation committed 1 year ago

1949691d

Merge pull request #22134 from gnecula:pallas_doc

a-googler committed 1 year ago

5fac179f

Improve error message when a global jax.Array is closed over a jitted function in McJAX.

yashk2810 committed 1 year ago

89c404e7

Fix error in custom call registration for some FFI functions

gnecula committed 1 year ago

2f808e9d

Update XLA dependency to use revision

Google-ML-Automation committed 1 year ago

2e83c324

[pallas] Added more documentation for grid and BlockSpec.

gnecula committed 1 year ago

bfdf8f4b

Merge pull request #22187 from selamw1:array_equal_doc

a-googler committed 1 year ago

071da567

array_equal_doscstring_added

selamw1 committed 1 year ago

74d59019

Simplify the ignore key when trace_context_in_key is False for `util.cache`

yashk2810 committed 1 year ago

75e7172e

Merge pull request #22180 from jakevdp:api-compat

a-googler committed 1 year ago

167f68bc

Explicitly disallow duplicated devices during array construction

junwhanahn committed 1 year ago

1a91fe76

[jax] Add a test that runs reduction on host

ezhulenev committed 1 year ago

de633956

doc: remove references to submodules that no longer exist

jakevdp committed 1 year ago

671db54f

Deprecate jax.interpreters xb, xc, xe abbreviations.

jakevdp committed 1 year ago

251dfcad

Reverts 56526b46aad726ec4632fc3b18c79d26a8f399ef

Google-ML-Automation committed 1 year ago

8c889b50

Update XLA dependency to use revision

Google-ML-Automation committed 1 year ago

a256b44e

Update C++ registration of `cu_lu_pivots_to_permutation` to use `XLA_FFI_REGISTER_HANDLER`

dfm committed 1 year ago

9b33df64

Make `device_put` work with inputs which are host local and the sharding is global sharding i.e. sharding spanning across multiple hosts.

yashk2810 committed 1 year ago

061f4df8

remove the cloud TPU disable of memories_test.py because everything should work now

yashk2810 committed 1 year ago

ba88601b

Rollback of #22157 due to internal breakage

Jake VanderPlas committed 1 year ago

56526b46

Make lowering aware of `compute_type` so that we choose the correct lowering code.

yashk2810 committed 1 year ago

ba5b3c79

Finalize deprecation of several previously-deprecated jax.core functions:

Jake VanderPlas committed 1 year ago

fbcb157a

Be stricter when it comes to handling dtypes in splash_attention mask function

apaszke committed 1 year ago

648b9519

Merge pull request #22159 from jakevdp:arange-doc

a-googler committed 1 year ago

44071f85

[Mosaic GPU] Add support for non-128B swizzles in WGMMA

apaszke committed 1 year ago

b19ad5b3

Ported threefry2x32 for GPU to the typed XLA FFI

gnecula committed 1 year ago

cbe52429

[Mosaic GPU] Stop using nvgpu for TMA

apaszke committed 1 year ago

3ebebdfb

[Mosaic GPU] Avoid failing when importing profiler.py even if lib.mosaic_gpu is unavailable.

Google-ML-Automation committed 1 year ago

ad4c9ab8

[Mosaic] Fix bug in VectorLayout::generalizes after cl/647395486

tlongeri committed 1 year ago

3a21c81e

Disable pallas_pipeline_tpu_test under asan/msan/tsan.

apaszke committed 1 year ago

bdefb812

Newer Older