ROCm/jax

Pull Requests Commits

Account for versioned clang binaries

charleshofer committed 232 days ago

83fa8c5d

Added gfx950 platform support, dropped unsupported platform

zahiqbal committed 236 days ago

a02fdd28

JAX release v0.6.0

hawkinsp committed 237 days ago

9e9a3f03

[Pallas] Fix relative links in pipelining docs

justinjfu committed 237 days ago

28538e4c

Fix bug in argnums_partial_except when static_argnums is unsorted.

dfm committed 237 days ago

34234010

Return False in `is_env_present` if importing kubernetes leads to a ModuleNotFoundError

yashk2810 committed 237 days ago

8d3a92e6

Fix a bug in reduce_window sharding rule where padding is a tuple but we were checking for a scalar instead. Fixes https://github.com/jax-ml/jax/issues/28070

yashk2810 committed 237 days ago

33d43f7a

[Pallas] Move pipelining doc figures to static folder

justinjfu committed 237 days ago

a9465e70

Add explicit sharding to Introduction to parallel programming doc

yashk2810 committed 237 days ago

24fdc0ad

Update XLA dependency to use revision

Google-ML-Automation committed 237 days ago

6b94a5ce

[Pallas] Refactor pipelining docs.

justinjfu committed 237 days ago

17e2dcaf

Update job name to show "build only" on Linux ARM64.

nitins17 committed 237 days ago

bc4fcd87

disable packed layout test on old arch prior to Hopper

Cjkkkk committed 237 days ago

fb875101

Update XLA dependency to use revision

Google-ML-Automation committed 237 days ago

a5236392

Fix race suppression for 3.13

vfdev-5 committed 237 days ago

8d474a3f

Add basic documentation for Pallas:MGPU layouts and transforms

apaszke committed 237 days ago

273e6efc

[pallas] Use `*Params` dataclasses to pass `compiler_params=` to `pl.pallas_call`

superbobry committed 237 days ago

bb6ca370

[Pallas:MGPU] Add a diagram outlining GPU memory spaces

apaszke committed 237 days ago

b7a041cc

Add docs for GPU memory spaces in Pallas:MGPU

apaszke committed 237 days ago

fe96e029

[Pallas][Mosaic][TPU] Add `disable_bounds_checks` compiler params

bythew3i committed 237 days ago

3a54f93d

Fix test flakyness by blocking until the data is ready.

pschuh committed 237 days ago

3ad33334

[pallas] Fix case of `Fusible{ElementDtype,TyRules}`.

chr1sj0nes committed 237 days ago

3dfbdd09

remove inaccurate inline comment in `PRNGKeyArray` constructor

froystig committed 237 days ago

be3c4d70

convert NumPy RNG key data to uncommitted default-device-backed `jax.Array` data

froystig committed 237 days ago

b35f5f7d

Minor improvements to doc for jax.nn.logsumexp.

carlosgmartin committed 237 days ago

b13364fb

Clean up softmax initial deprecation

jakevdp committed 237 days ago

1aeac03c

Enable standard_insert_pvary for optimization_barrier which was disabled before.

yashk2810 committed 237 days ago

c3d8ac79

[ragged-paged-attn][NFC] Set kv_pages_per_blk uplimit.

bythew3i committed 237 days ago

5214342e

Fix bugs in tp_traverse handlers.

hawkinsp committed 237 days ago

0f2b3f21

Roll back https://github.com/jax-ml/jax/pull/28022 due to test breakages.

Jake VanderPlas committed 237 days ago

c6a8050c

Older