ROCm/jax

Pull Requests Commits

[ROCM] Fix version checks after rocm pjrt integration

Ruturaj4 committed 2 years ago

637370ba

Run single controller device_put via efficient reshard if the device_set of input and the sharding is the same. The transfer_guard in the test fails before this CL.

yashk2810 committed 2 years ago

94ba6c3f

[pallas] Add limited support for shape polymorphism for TPU

gnecula committed 2 years ago

9def0f1c

Don't use the trace context in the prune_closed_jaxpr_outputs cache.

hawkinsp committed 2 years ago

55589fbf

[Pallas] Add vmap/batched key support to to_pallas_key. This is helpful for workflows where the key is split before being passed into the kernel.

justinjfu committed 2 years ago

9653f58f

Merge pull request #22132 from rajasekharporeddy:testbranch2

a-googler committed 2 years ago

7cd5b5c6

Merge pull request #22208 from superbobry:docs

a-googler committed 2 years ago

32a70718

Revived the workaround for not-expanding type aliases

superbobry committed 2 years ago

e80632e6

Avoid different return value when running on single host.

andsteing committed 2 years ago

a386af44

Bump up the shard_count for GPU FFT tests

apaszke committed 2 years ago

727d1204

Update jnp.mean doc

rajasekharporeddy committed 2 years ago

9623dcfa

Update XLA dependency to use revision

Google-ML-Automation committed 2 years ago

1949691d

Merge pull request #22134 from gnecula:pallas_doc

a-googler committed 2 years ago

5fac179f

Improve error message when a global jax.Array is closed over a jitted function in McJAX.

yashk2810 committed 2 years ago

89c404e7

Fix error in custom call registration for some FFI functions

gnecula committed 2 years ago

2f808e9d

Update XLA dependency to use revision

Google-ML-Automation committed 2 years ago

2e83c324

[pallas] Added more documentation for grid and BlockSpec.

gnecula committed 2 years ago

bfdf8f4b

Merge pull request #22187 from selamw1:array_equal_doc

a-googler committed 2 years ago

071da567

array_equal_doscstring_added

selamw1 committed 2 years ago

74d59019

Simplify the ignore key when trace_context_in_key is False for `util.cache`

yashk2810 committed 2 years ago

75e7172e

Merge pull request #22180 from jakevdp:api-compat

a-googler committed 2 years ago

167f68bc

Explicitly disallow duplicated devices during array construction

junwhanahn committed 2 years ago

1a91fe76

[jax] Add a test that runs reduction on host

ezhulenev committed 2 years ago

de633956

doc: remove references to submodules that no longer exist

jakevdp committed 2 years ago

671db54f

Reverts 56526b46aad726ec4632fc3b18c79d26a8f399ef

Google-ML-Automation committed 2 years ago

8c889b50

Update XLA dependency to use revision

Google-ML-Automation committed 2 years ago

a256b44e

Update C++ registration of `cu_lu_pivots_to_permutation` to use `XLA_FFI_REGISTER_HANDLER`

dfm committed 2 years ago

9b33df64

Make `device_put` work with inputs which are host local and the sharding is global sharding i.e. sharding spanning across multiple hosts.

yashk2810 committed 2 years ago

061f4df8

remove the cloud TPU disable of memories_test.py because everything should work now

yashk2810 committed 2 years ago

ba88601b

Rollback of #22157 due to internal breakage

Jake VanderPlas committed 2 years ago

56526b46

Older