ROCm/jax

Pull Requests Commits

[ROCm] Add f8 test for ROCm

Ruturaj4 committed 1 year ago

6cd36d6b

Update XLA dependency to use revision

a-googler committed 1 year ago

b1cb90cc

Merge pull request #20195 from Micky774:array_api_astype

a-googler committed 1 year ago

493698e6

Merge pull request #20317 from inailuig:mpi_collectives

a-googler committed 1 year ago

d20a2f10

Rename arg in build script to be more clear.

Jieying Luo committed 1 year ago

16b4f697

Return early from eigh for small matrices. This was accidentally removed cl/577222219.

a-googler committed 1 year ago

47c9495d

Merge pull request #20849 from mattjj:jit-docstring-tweaks

a-googler committed 1 year ago

1f4c31d0

Support auto in shard_map.

pschuh committed 1 year ago

7ba811eb

Merge pull request #20870 from selamw1:jax_doc

a-googler committed 1 year ago

a8ee946e

Merge pull request #20867 from google:dependabot/github_actions/actions/download-artifact-4.1.6

a-googler committed 1 year ago

461c5ec4

Merge pull request #20868 from google:dependabot/github_actions/actions/upload-artifact-4.3.3

a-googler committed 1 year ago

1bb82525

redundant phrase 'ever time' removed

selamw1 committed 1 year ago

b02f82b8

Bump actions/upload-artifact from 4.3.1 to 4.3.3

dependabot[bot] committed 1 year ago

Verified e93954c1

Bump actions/download-artifact from 4.1.4 to 4.1.6

dependabot[bot] committed 1 year ago

Verified 4b79d362

Optimize `jax.device_put()` dispatch for 1:1 device-to-device transfers

junwhanahn committed 1 year ago

4be25d71

Enable some more C API tests.

Jieying Luo committed 1 year ago

1b1c6e7c

Add support for copy kwarg in astype to match Array API

Micky774 committed 1 year ago

30cd3b88

Add some docstrings for remote DMAs and semaphore barriers.

a-googler committed 1 year ago

667a0c1f

inailuig committed 1 year ago

1e32fb51

[Mosaic:GPU] Update lowering to match upstream changes in the LLVM dialect

apaszke committed 1 year ago

b79f3b77

Update XLA dependency to use revision

a-googler committed 1 year ago

83aff78d

[XLA:TPU] Support scan loops for parameter input and output streaming in host offloading.

jvstokes committed 1 year ago

fd100780

Update XLA dependency to use revision

a-googler committed 1 year ago

e498bca2

Merge some loops in device_put since it's trivial to do so

yashk2810 committed 1 year ago

1837b436

Update XLA dependency to use revision

a-googler committed 1 year ago

0943eb38

tweak jit docstring

mattjj committed 1 year ago

b8df23c2

Removed unnecessary jax.tree.map calls from *_callback_impl functions

superbobry committed 1 year ago

5981df7b

try remove MpiCollectives from type annotation

inailuig committed 1 year ago

fdd24d13

Merge pull request #20840 from sagelywizard:patch-2

a-googler committed 1 year ago

52f5f703

Merge pull request #20836 from mattjj:quickstart-tweaks

a-googler committed 1 year ago

84248249

Older