ROCm/jax

Pull Requests Commits

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

blakehechtman committed 300 days ago

60212a39

Add supports_pinned_allocator to allow debugging pinning issues.

pschuh committed 300 days ago

f8a3f061

Add host offloading docs to public website

yashk2810 committed 300 days ago

56bbaa36

Fix empty string handling for cloud_tpu_cluster

Google-ML-Automation committed 300 days ago

da8f62b4

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

yashk2810 committed 300 days ago

a9c49ac0

Integrate LLVM at llvm/llvm-project@2d287f51eff2

alinas committed 300 days ago

8683b76d

Remove type annotation of get_gpu_client

hhb committed 301 days ago

4786d122

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

matthiaskramm committed 301 days ago

3712656b

Allow unreduced propagation only for `add` right now.

yashk2810 committed 301 days ago

d2284bf8

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

justinjfu committed 301 days ago

44a30b20

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

mwhittaker committed 301 days ago

515f81bf

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

naummo committed 301 days ago

6c66e938

[jaxlib] Add compile_and_load, compile_and_load_ifrt_program to xla_client stub.

danielsuo committed 301 days ago

08385f51

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

yashk2810 committed 301 days ago

098624dc

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

Google-ML-Automation committed 301 days ago

5d0cd68e

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

justinjfu committed 301 days ago

f15ad5a3

Merge pull request #28610 from jakevdp:fix-put

Google-ML-Automation committed 301 days ago

4fbd7f1d

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

justinjfu committed 301 days ago

bf830823

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

yashk2810 committed 301 days ago

93d02349

Merge pull request #28279 from gspschmid:gschmid/jax_fwd_and_bwd_argnums

Google-ML-Automation committed 301 days ago

87d16157

Fix handling of final style primitives in pallas cost estimate.

dfm committed 301 days ago

8c987bfb

[Pallas][Mosaic GPU] Expand TMEM support.

justinjfu committed 301 days ago

f7e11660

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

nitins17 committed 301 days ago

66fdee2d

jnp.linalg.matrix_power: support non-float inputs

jakevdp committed 301 days ago

1e1f1e0d

jnp.put: check inplace before other conditions

jakevdp committed 301 days ago

4013965d

Record event start time in `dispatch.LogElapsedTimeContextManager`.

danielsuo committed 301 days ago

5c77f249

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

danielsuo committed 301 days ago

9f64ddd3

Update XLA dependency to use revision

Google-ML-Automation committed 301 days ago

dbfc93fd

Merge pull request #28606 from dfm:fix-stemmer

Google-ML-Automation committed 301 days ago

0341b6f3

Fix docs build by constraining snowballstemmer version.

dfm committed 301 days ago

93eb725e

Older