ROCm/jax

Pull Requests Commits

[JAX:ANN] Make a test more reliable by increasing the sample size

blakehechtman committed 1 year ago

60212a39

Add supports_pinned_allocator to allow debugging pinning issues.

pschuh committed 1 year ago

f8a3f061

Add host offloading docs to public website

yashk2810 committed 1 year ago

56bbaa36

Fix empty string handling for cloud_tpu_cluster

Google-ML-Automation committed 1 year ago

da8f62b4

Simplify add's unreduced rule. Only propagate unreduced if both lhs and rhs are unreduced.

yashk2810 committed 1 year ago

a9c49ac0

Integrate LLVM at llvm/llvm-project@2d287f51eff2

alinas committed 1 year ago

8683b76d

Remove type annotation of get_gpu_client

hhb committed 1 year ago

4786d122

Fix the PyTest TPU jobs on the Continuous Wheel Tests workflow.

matthiaskramm committed 1 year ago

3712656b

Allow unreduced propagation only for `add` right now.

yashk2810 committed 1 year ago

d2284bf8

[Pallas][Mosaic GPU] Add transpose support to tcgen05_mma

justinjfu committed 1 year ago

44a30b20

Shut down `PreemptionSyncManager` when `jax.distributed.shutdown()` is called.

mwhittaker committed 1 year ago

515f81bf

[Mosaic][SC] Expose control over the number of active SC cores

naummo committed 1 year ago

6c66e938

[jaxlib] Add compile_and_load, compile_and_load_ifrt_program to xla_client stub.

danielsuo committed 1 year ago

08385f51

Make the type checker match the runtime behavior of PartitionSpec not inherting from a tuple.

yashk2810 committed 1 year ago

098624dc

Merge pull request #28612 from jakevdp:fix-matrix-power

Google-ML-Automation committed 1 year ago

5d0cd68e

[Pallas][Mosaic GPU] Refactor carry_coroutine in warp specialized pipeline to be a callback instead of a coroutine.

justinjfu committed 1 year ago

f15ad5a3

Merge pull request #28610 from jakevdp:fix-put

Google-ML-Automation committed 1 year ago

4fbd7f1d

[Pallas][Mosaic GPU] Add collective support to Blackwell/tcgen05 MMA.

justinjfu committed 1 year ago

bf830823

Add `out_sharding` to `jnp.repeat`. Drop into auto mode if out_sharding is provided.

yashk2810 committed 1 year ago

93d02349

Merge pull request #28279 from gspschmid:gschmid/jax_fwd_and_bwd_argnums

Google-ML-Automation committed 1 year ago

87d16157

Fix handling of final style primitives in pallas cost estimate.

dfm committed 1 year ago

8c987bfb

[Pallas][Mosaic GPU] Expand TMEM support.

justinjfu committed 1 year ago

f7e11660

Upgrade Mac CI builds to run on the Sequoia pool and Apple Clang 17

nitins17 committed 1 year ago

66fdee2d

jnp.linalg.matrix_power: support non-float inputs

jakevdp committed 1 year ago

1e1f1e0d

jnp.put: check inplace before other conditions

jakevdp committed 1 year ago

4013965d

Record event start time in `dispatch.LogElapsedTimeContextManager`.

danielsuo committed 1 year ago

5c77f249

[xla::PyClient] Update PyClient to use xla::ifrt::CompileAndLoad.

danielsuo committed 1 year ago

9f64ddd3

Update XLA dependency to use revision

Google-ML-Automation committed 1 year ago

dbfc93fd

Merge pull request #28606 from dfm:fix-stemmer

Google-ML-Automation committed 1 year ago

0341b6f3

Fix docs build by constraining snowballstemmer version.

dfm committed 1 year ago

93eb725e

Older