ROCm/jax

Pull Requests Commits

Rollback because of breakages

Google-ML-Automation committed 1 year ago

411928b9

Fix and reenable cudnn_fusion_test.

sergachev committed 1 year ago

b320dc2e

[pallas:mosaic_gpu] Do not DCE the jaxpr in the lowering pass

superbobry committed 1 year ago

b3fca904

[Pallas/MGPU] Undo transforms on refs before giving them back to the users

apaszke committed 1 year ago

21fea5b0

Exposed sequential iteration index via `pl.program_id` in Pallas Mosaic GPU

superbobry committed 1 year ago

38d2a573

Update XLA dependency to use revision

Google-ML-Automation committed 1 year ago

2cfbdb6c

Update docs for jnp.trapezoid

jakevdp committed 1 year ago

7cc5ea4c

Update XLA dependency to use revision

Google-ML-Automation committed 1 year ago

6790b90f

Merge pull request #23982 from dfm:ffi-call-effects

Google-ML-Automation committed 1 year ago

15024baa

Add support for FFI calls with side effects via ffi_call

dfm committed 1 year ago

d80a89d8

Bump test tolerance on FFT test that started failing in CI after an XLA change.

hawkinsp committed 1 year ago

061f435b

Fix test failure when shardy is not enabled.

hawkinsp committed 1 year ago

366c8238

Fix tests that ask for an accelerator but don't use it.

hawkinsp committed 1 year ago

5969e799

Merge pull request #23957 from jakevdp:choice-doc

Google-ML-Automation committed 1 year ago

ff0a98a2

Merge pull request #23981 from jakevdp:checkout-pin

Google-ML-Automation committed 1 year ago

4e8a763c

Merge pull request #23977 from rajasekharporeddy:testbranch1

Google-ML-Automation committed 1 year ago

2963cbed

actions/checkout: pin specific hash

jakevdp committed 1 year ago

f0e27fd7

Better documentation for jnp.choose

jakevdp committed 1 year ago

c0612576

Update XLA dependency to use revision

Google-ML-Automation committed 1 year ago

20122ff7

Improve docs for jax.numpy: arcsinh, arccosh and arctanh

rajasekharporeddy committed 1 year ago

c17ae0fd

Merge pull request #23870 from Zantares:tenglu/flush_output

Google-ML-Automation committed 1 year ago

df042fde

Merge pull request #23804 from hawkinsp:vis

Google-ML-Automation committed 1 year ago

a6167585

Merge pull request #23961 from apivovarov:pylint2

Google-ML-Automation committed 1 year ago

af12e0a1

Merge pull request #23980 from jakevdp:update-actions

Google-ML-Automation committed 1 year ago

898f0572

Merge pull request #23965 from zhenying-liu:weight-offloading-test

Google-ML-Automation committed 1 year ago

b7622911

actions: remove use of ratchet tool

jakevdp committed 1 year ago

cb9f26c8

Include the sdy MLIR dialect in jaxlib.

hawkinsp committed 1 year ago

5a1d0a6c

Don't pass --nocheck_visibility to Bazel.

hawkinsp committed 1 year ago

e4790b63

Replace disable_backends with enable_backends on jax_multiplatform_test.

hawkinsp committed 1 year ago

26632fd3

[Pallas/MGPU] Skip output transfers when they don't depend on sequenital dims

apaszke committed 1 year ago

5740ab3b

Newer Older