ROCm/jax

Pull Requests Commits

Fix file path regexp for Windows

gnecula committed 1 year ago

c1734f22

Update RNN kernels to use FFI.

dfm committed 1 year ago

c6e83903

Merge pull request #26380 from dfm:remove-legacy-custom-call-doc

Google-ML-Automation committed 1 year ago

cce3df10

Remove the tutorial about legacy custom calls.

dfm committed 1 year ago

21281f37

Merge pull request #26375 from jakevdp:fix-doc

Google-ML-Automation committed 1 year ago

ce51384b

Merge pull request #26372 from jakevdp:split-lax-numpy

Google-ML-Automation committed 1 year ago

af29a7b0

DOC: update numpy/scipy versions in deprecation doc

jakevdp committed 1 year ago

e0ee4156

refactor: move scalar type defs out of lax_numpy.py

jakevdp committed 1 year ago

b4f98eef

Default JAX_CPU_COLLECTIVES_IMPLEMENTATION to 'gloo'.

skye committed 1 year ago

f07243a7

Merge pull request #25097 from jburnim:jburnim_pallas_interpret_mode

Google-ML-Automation committed 1 year ago

4b86ff22

Remove code for compatibility with jaxlib < 0.5.0.

hawkinsp committed 1 year ago

840192d3

Start a new TPU interpret mode for Pallas.

jburnim committed 1 year ago

1c82484c

Merge pull request #26369 from justinjfu:sourcemap_win

Google-ML-Automation committed 1 year ago

1b7b04f7

Temporarily disable source mapper on win32

justinjfu committed 1 year ago

49fa1274

Update the sparse GPU kernels in jaxlib to use the FFI.

dfm committed 1 year ago

5e915d33

Merge pull request #26342 from jakevdp:lax-docs

Google-ML-Automation committed 1 year ago

f54ddb31

doc: improve docs for jax.lax trig functions

jakevdp committed 1 year ago

2fb750e0

Merge pull request #26363 from mattjj:mutable-array-error-tweak

Google-ML-Automation committed 1 year ago

7025b7d1

[Mosaic GPU] Add more SMEM buffers to avoid blocking for memory traffic

apaszke committed 1 year ago

6524e67f

[Mosaic GPU] Relax TMEM stride constraints on dimensions of size 1

apaszke committed 1 year ago

a61d8002

[Pallas TPU] Remove obsolete skip condition

ayaka14732 committed 1 year ago

afad924d

tweak an error message on mutable arrays

mattjj committed 1 year ago

96b6f77c

Increase test shard_count for shape_poly_test on GPU

gnecula committed 1 year ago

a678396f

[Mosaic GPU] Take TMEM as a TMEMRef in tcgen05.mma, not as a raw address

apaszke committed 1 year ago

026b6c97

[pallas:triton] Temporarily reverted to the lowering using Triton IR

superbobry committed 1 year ago

efbb0afd

Merge pull request #26348 from gnecula:debug_info_jaxpr_3

Google-ML-Automation committed 1 year ago

5d647ccf

Merge pull request #26279 from MichaelHudgins:tsan-resultstore

MichaelHudgins committed 1 year ago

2e808f28

[better_errors] Continue adding debug info to Jaxprs (step 3)

gnecula committed 1 year ago

904b7486

Update tsan job to upload to resultstore

MichaelHudgins committed 1 year ago

be0f3d8f

If sharding is not None (that's passed to convert_element_type), only compare it with operand's sharding if the sharding is concrete. Otherwise doing `getattr(operand, 'sharding')` on a `Tracer` leads to weird timeouts.

yashk2810 committed 1 year ago

0fb278a0

Older