ROCm/jax

Pull Requests Commits

Split JaxTestLoader and related classes into a separate file.

hawkinsp committed 1 year ago

382285d3

Merge pull request #27895 from jakevdp:random-test-refactor

Google-ML-Automation committed 1 year ago

cf268a7f

Restructure WeakrefLRUCache.

hawkinsp committed 1 year ago

2ceb97cd

Fix build breakage from missing _sdy_enums_gen.py.

hawkinsp committed 1 year ago

a5055206

Make printing work with shard_map after vma has been switched on

yashk2810 committed 1 year ago

93202140

Refactor random_lax_test.py

jakevdp committed 1 year ago

c5bd13bb

[JAX] Remove deprecated exports in jax.lib.xla_client.

hawkinsp committed 1 year ago

713ea3ca

Merge pull request #27888 from vfdev-5:add-314-cpython-update_one_slot-suppress

Google-ML-Automation committed 1 year ago

178b2f13

Set `jax_varying_axes_in_types` to True by default.

yashk2810 committed 1 year ago

75e4279e

Merge pull request #27893 from mattjj:shmap-vma-docs

Google-ML-Automation committed 1 year ago

fbca0909

[shard-map] add docs for VMAs

mattjj committed 1 year ago

97a0d752

Remove xla_extension.stablehlo_to_mhlo.

hawkinsp committed 1 year ago

56646afa

Skip jax/tests:unary_ops_accuracy_test when running with older versions of StableHLO.

GleasonK committed 1 year ago

7c1595ac

Added update_one_slot race tsan suppression for 3.14

vfdev-5 committed 1 year ago

e9308a21

Fix docstrings of segment_{prod,max,min} after commit 4679f45.

arnoegw committed 1 year ago

e772a08a

Merge pull request #27886 from mattjj:26763

Google-ML-Automation committed 1 year ago

c418495b

Merge pull request #27759 from mattjj:vmappable-bind-fix

Google-ML-Automation committed 1 year ago

2863b488

[shard-map] make shard_map work with custom_jvp symbolic zeros

mattjj committed 1 year ago

2b3839d2

Merge pull request #27853 from carlosgmartin:merge_tuple_update_tuple_replace

Google-ML-Automation committed 1 year ago

03856671

Add option for debug print to be called on partitioned arguments rather than all-gather and print the entire logical arguments.

danielsuo committed 1 year ago

e750d7ea

Merge pull request #27879 from mattjj:shmap-fix-5

Google-ML-Automation committed 1 year ago

21a4429a

Merge tuple_replace and tuple_update in jax._src.util.

carlosgmartin committed 1 year ago

b6a46310

[shard-map] fix another bug where we incorrectly handled None in check_rep

mattjj committed 1 year ago

8383af01

More changes for enabling `vma` by default in JAX

yashk2810 committed 1 year ago

76c6b5b0

Update XLA dependency to use revision

Google-ML-Automation committed 1 year ago

5f4c6e49

[Mosaic GPU] Add a `LayoutCast` op to the Mosaic GPU mlir dialect.

dimitar-asenov committed 1 year ago

9adc3cc0

Merge pull request #27870 from medaminezghal:fix-clang19

Google-ML-Automation committed 1 year ago

8521c017

Fix build for clang >=19

medaminezghal committed 1 year ago

3d006e9f

Update XLA dependency to use revision

Google-ML-Automation committed 1 year ago

d9d2f0b0

Merge pull request #27866 from hawkinsp:race

Google-ML-Automation committed 1 year ago

90cf9b4c

Older