ROCm/jax

Pull Requests Commits

[ROCm]: Add support to continue on fail

Rahul Batra committed 2 years ago

e04ccdfe

Merge pull request #20802 from jakevdp:dup-thinking-in-jax

a-googler committed 2 years ago

1c8534e0

Merge pull request #20800 from jakevdp:pin-sphinx

a-googler committed 2 years ago

73bd54f0

DOC: remove copy of thinking-in-jax from new tutorial flow

jakevdp committed 2 years ago

edfdc36b

DOC: pin sphinx to >=7.3.2

jakevdp committed 2 years ago

29b15e2e

Fix Pallas' while_loop lowering to properly account for the loop length

apaszke committed 2 years ago

acde885b

Reverts f18739900c615a85e8d182bcf3217f704cf7aa0d

a-googler committed 2 years ago

6963c770

Update XLA dependency to use revision

a-googler committed 2 years ago

00e3a1ea

Merge pull request #20792 from jakevdp:fix-sharding-doc

a-googler committed 2 years ago

1d615a3b

Merge pull request #20732 from jakevdp:doc-stateful

a-googler committed 2 years ago

fc654a57

Merge pull request #20794 from jakevdp:pin-sphinx

a-googler committed 2 years ago

07063fb1

Merge pull request #20784 from jakevdp:doc-installation

a-googler committed 2 years ago

e554416d

DOC: add stateful computation doc

jakevdp committed 2 years ago

3c18a021

Pin sphinx version to avoid error in 7.3.0

jakevdp committed 2 years ago

40ad5b28

DOC: fix Mesh construction in sharding doc

jakevdp committed 2 years ago

2552c290

Merge pull request #20787 from superbobry:lazy-imports

a-googler committed 2 years ago

1c4195c5

Import rich lazily

superbobry committed 2 years ago

1be54511

Import etils.epath lazily.

hawkinsp committed 2 years ago

f1873990

Temporarily disable async dispatch on JAX CPU by setting 'jax_cpu_enable_async_dispatch' to be `False` by default, as we observed abnormal memory usage increases.

yueshengys committed 2 years ago

1f83908b

Merge pull request #20756 from Micky774:array-api-cumulative-sum

a-googler committed 2 years ago

47815c54

Add new cumulative_sum function to numpy and array_api

Micky774 committed 2 years ago

ceeb9757

DOC: update installation guide

jakevdp committed 2 years ago

a4ccb354

Merge pull request #20778 from superbobry:lazy-imports

a-googler committed 2 years ago

adbb11f9

[Mosaic] Support scf.while and scf.condition.

a-googler committed 2 years ago

5bd6013e

Update JAX official doc: point out that the device numbers are not in numerical order because of the underlying torus hardware topology.

yueshengys committed 2 years ago

1a650cdc

Merge pull request #20779 from jakevdp:sharding-101

a-googler committed 2 years ago

e83b0ce3

Merge pull request #20773 from jakevdp:update-array-api

a-googler committed 2 years ago

42398a54

Merge pull request #20657 from Sai-Suraj-27:fix_error_messages

a-googler committed 2 years ago

d2a2b9f2

DOC: add introduction to sharded computation

jakevdp committed 2 years ago

8e34da70

Import web_pdb lazily

superbobry committed 2 years ago

35136851

Older