ROCm/jax

Pull Requests Commits

Pallas bitwise_left_shift unit test fix

Rahul Batra committed 2 years ago

d2a622dd

Merge pull request #21763 from gnecula:export_api

a-googler committed 2 years ago

3fe73777

[export] Create the jax.export module APIs.

gnecula committed 2 years ago

b33aca6b

Allow generation of sharding strategies with mixed mesh shapes by default.

Google-ML-Automation committed 2 years ago

833c7ba7

[Mosaic GPU] Don't always run with llvm::DebugFlag enabled

apaszke committed 2 years ago

0739d520

Add initialization annotations (for the benefit of MSAN) to variables that are initialized by external functions.

tkoeppe committed 2 years ago

cd93b46d

Merge pull request #21765 from hawkinsp:release

a-googler committed 2 years ago

991797a8

Removed the double re-exporting of Pallas GPU/TPU APIs

superbobry committed 2 years ago

5e7ad600

[Mosaic GPU] Load LLVM lowering interfaces for all dialects

apaszke committed 2 years ago

3b4039c8

Prepare for 0.4.29 release.

hawkinsp committed 2 years ago

e071053d

[pallas] Move the hardware_generation query in the code path that needs it

gnecula committed 2 years ago

2ade7e75

Merge pull request #21759 from rajasekharporeddy:testbranch1

a-googler committed 2 years ago

af95803d

Fix Typos in docs and one error message

rajasekharporeddy committed 2 years ago

775c6f87

Update XLA dependency to use revision

Google-ML-Automation committed 2 years ago

8fbe65b4

Expand `device_put` benchmarks to run with different numbers of arrays and input types

junwhanahn committed 2 years ago

6617a0d1

Issue a warning where code relies on a bug where treedef.flatten_up_to(...) was overly permissive for None treedefs.

hawkinsp committed 2 years ago

a8246ea6

[export] Move the export implementation to jax._src.export.

gnecula committed 2 years ago

14d87d3b

Update XLA dependency to use revision

Google-ML-Automation committed 2 years ago

aaa559a6

Merge pull request #21507 from renecotyfanboy:main

a-googler committed 2 years ago

b486a951

Update XLA dependency to use revision

Google-ML-Automation committed 2 years ago

6c822c01

Avoid "min() arg is an empty sequence" error after enabling "jax_explain_cache_misses".

Google-ML-Automation committed 2 years ago

d3240402

increase default precision for hyp1f1

renecotyfanboy committed 2 years ago

751d59ce

Add a no input memories_test and enable memories test on vf 2x2

yashk2810 committed 2 years ago

57826d8c

Merge pull request #21718 from jakevdp:pallas-config

a-googler committed 2 years ago

0d047a11

Merge code between `make_jaxpr` and `jit(f).trace`.

yashk2810 committed 2 years ago

44a13c9d

Merge pull request #21615 from selamw1:append_doc

a-googler committed 2 years ago

25cc84b8

Merge pull request #21744 from superbobry:typing

a-googler committed 2 years ago

dfc6076d

Added filelock to py_deps

superbobry committed 2 years ago

136289e9

Merge pull request #21298 from oliverdutton:pallas_interpreter_indexing_fix

a-googler committed 2 years ago

7d913f76

Removed unnecessary mypy exclusions from pyproject.toml

superbobry committed 2 years ago

0786da8f

Older