ROCm/jax

Pull Requests Commits

Merge remote-tracking branch 'origin/rocm-main' into ci-upstream-sync-55_1

invalid-email-address committed 1 year ago

20c5c714

Delete non-public API jax.lib.xla_bridge._backends

Jake VanderPlas committed 1 year ago

e6d6c4ef

[jax:custom_partitioning] Make SdyShardingRule a user facing class.

bixia1 committed 1 year ago

d4899f7b

Merge pull request #25357 from jakevdp:core-deps

Google-ML-Automation committed 1 year ago

9c32fe8f

[Pallas] Add non-square pl.dot test cases.

WindQAQ committed 1 year ago

e418e883

Deduplicate some GPU plugin definition code.

dfm committed 1 year ago

593143e1

Merge pull request #25378 from hawkinsp:optbarrier

Google-ML-Automation committed 1 year ago

210bd308

jax.core: deprecate a number of APIs

jakevdp committed 1 year ago

6541a620

Reverts bdadc53ebcd40a5091d66d2586deba82fe5e01ca

beckerhe committed 1 year ago

cb6881d9

[AutoPGLE] Share FDO profile even when compilation cache disabled.

Google-ML-Automation committed 1 year ago

e6dfe8f3

Fix mypy failure

Google-ML-Automation committed 1 year ago

6dbafed7

[AutoPGLE] Cleanup compiler code.

Google-ML-Automation committed 1 year ago

8813973d

Merge pull request #176 from ROCm/main

charleshofer committed 1 year ago

Verified f04f164a

Remove code in jax2tf for compatibility with TF 2.10 or earlier.

hawkinsp committed 1 year ago

acae2f05

Merge pull request #25369 from jax-ml:mutable-arrays-ad

Google-ML-Automation committed 1 year ago

263d4d14

Merge pull request #25374 from traversaro:patch-1

Google-ML-Automation committed 1 year ago

8e7aaa79

Update conda-forge installation docs after CUDA 12 upgrade

traversaro committed 1 year ago

Verified 09309e64

Merge pull request #25335 from gnecula:export_doc_call

Google-ML-Automation committed 1 year ago

90de28cd

Introduce `lax.ragged_all_to_all` primitive

ghpvnist committed 1 year ago

12c30578

Add a no-op batching rule for optimization_barrier_p

yashk2810 committed 1 year ago

944d822c

Merge pull request #25371 from jakevdp:mypy-numpy-version

Google-ML-Automation committed 1 year ago

1743f2c4

CI: temporarily pin numpy version for mypy check

jakevdp committed 1 year ago

a36af966

Add a `freeze` primitive to delimit ref lifetimes for AD.

dougalm committed 1 year ago

fc2edbfa

Merge pull request #25368 from hawkinsp:postrelease

Google-ML-Automation committed 1 year ago

7b5cb56f

Merge branch 'release/0.4.37' into main.

hawkinsp committed 1 year ago

820f51dc

Fix expected exception type in pallas grad tests.

dfm committed 1 year ago

978d35f6

Update XLA dependency to use revision

Google-ML-Automation committed 1 year ago

71c48cba

Port symmetric tridiagonal reduction GPU kernel to FFI.

dfm committed 1 year ago

32df37e6

Update versions for v0.4.37 release.

hawkinsp committed 1 year ago

ffb07cda

Use private names for args in api_util to avoid shadowing kwargs keys.

dougalm committed 1 year ago

95892fda

Older