ROCm/jax

Pull Requests Commits

[ROCM] make mosaic dependency cuda specific

Ruturaj4 committed 1 year ago

445a6923

[ROCM] add typed XLA FFI support in rocm specific code

Ruturaj4 committed 1 year ago

665d8c0b

[ROCM] Replace xla::Status with absl::Status in rocm_plugin_extension.cc

Ruturaj4 committed 1 year ago

fd6e1982

[ROCM] add pytest reportlog

Ruturaj4 committed 1 year ago

bd73f6bd

[JAX] Fix run_single_gpu.py to include correct paths

Ruturaj4 committed 1 year ago

2c082ffe

[ROCm]: Add support to continue on fail

Rahul Batra committed 1 year ago

6712b3de

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 1 year ago

94ba6c3f

[pallas] Add limited support for shape polymorphism for TPU

gnecula committed 1 year ago

9def0f1c

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

hawkinsp committed 1 year 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 1 year ago

9653f58f

Merge pull request #22132 from rajasekharporeddy:testbranch2

a-googler committed 1 year ago

7cd5b5c6

Merge pull request #22208 from superbobry:docs

a-googler committed 1 year ago

32a70718

Revived the workaround for not-expanding type aliases

superbobry committed 1 year ago

e80632e6

Avoid different return value when running on single host.

andsteing committed 1 year ago

a386af44

Bump up the shard_count for GPU FFT tests

apaszke committed 1 year ago

727d1204

Update jnp.mean doc

rajasekharporeddy committed 1 year ago

9623dcfa

Update XLA dependency to use revision

Google-ML-Automation committed 1 year ago

1949691d

Merge pull request #22134 from gnecula:pallas_doc

a-googler committed 1 year ago

5fac179f

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

yashk2810 committed 1 year ago

89c404e7

Fix error in custom call registration for some FFI functions

gnecula committed 1 year ago

2f808e9d

Update XLA dependency to use revision

Google-ML-Automation committed 1 year ago

2e83c324

[pallas] Added more documentation for grid and BlockSpec.

gnecula committed 1 year ago

bfdf8f4b

Merge pull request #22187 from selamw1:array_equal_doc

a-googler committed 1 year ago

071da567

array_equal_doscstring_added

selamw1 committed 1 year ago

74d59019

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

yashk2810 committed 1 year ago

75e7172e

Merge pull request #22180 from jakevdp:api-compat

a-googler committed 1 year ago

167f68bc

Explicitly disallow duplicated devices during array construction

junwhanahn committed 1 year ago

1a91fe76

[jax] Add a test that runs reduction on host

ezhulenev committed 1 year ago

de633956

doc: remove references to submodules that no longer exist

jakevdp committed 1 year ago

671db54f

Reverts 56526b46aad726ec4632fc3b18c79d26a8f399ef

Google-ML-Automation committed 1 year ago

8c889b50

Older