ROCm/jax

Pull Requests Commits

Revised supported gfx-arch, fixed build script

zahiqbal committed 290 days ago

88c9b5dc

Fixes for ci build

mrodden committed 290 days ago

a7833b89

JAX release v0.6.0

hawkinsp committed 290 days ago

9e9a3f03

[Pallas] Fix relative links in pipelining docs

justinjfu committed 290 days ago

28538e4c

Fix bug in argnums_partial_except when static_argnums is unsorted.

dfm committed 290 days ago

34234010

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

yashk2810 committed 290 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 290 days ago

33d43f7a

[Pallas] Move pipelining doc figures to static folder

justinjfu committed 290 days ago

a9465e70

Add explicit sharding to Introduction to parallel programming doc

yashk2810 committed 290 days ago

24fdc0ad

Update XLA dependency to use revision

Google-ML-Automation committed 290 days ago

6b94a5ce

[Pallas] Refactor pipelining docs.

justinjfu committed 290 days ago

17e2dcaf

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

nitins17 committed 290 days ago

bc4fcd87

disable packed layout test on old arch prior to Hopper

Cjkkkk committed 290 days ago

fb875101

Update XLA dependency to use revision

Google-ML-Automation committed 290 days ago

a5236392

Fix race suppression for 3.13

vfdev-5 committed 290 days ago

8d474a3f

Add basic documentation for Pallas:MGPU layouts and transforms

apaszke committed 290 days ago

273e6efc

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

superbobry committed 290 days ago

bb6ca370

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

apaszke committed 290 days ago

b7a041cc

Add docs for GPU memory spaces in Pallas:MGPU

apaszke committed 290 days ago

fe96e029

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

bythew3i committed 290 days ago

3a54f93d

Fix test flakyness by blocking until the data is ready.

pschuh committed 290 days ago

3ad33334

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

chr1sj0nes committed 290 days ago

3dfbdd09

remove inaccurate inline comment in `PRNGKeyArray` constructor

froystig committed 290 days ago

be3c4d70

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

froystig committed 290 days ago

b35f5f7d

Minor improvements to doc for jax.nn.logsumexp.

carlosgmartin committed 290 days ago

b13364fb

Clean up softmax initial deprecation

jakevdp committed 290 days ago

1aeac03c

Enable standard_insert_pvary for optimization_barrier which was disabled before.

yashk2810 committed 290 days ago

c3d8ac79

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

bythew3i committed 290 days ago

5214342e

Fix bugs in tp_traverse handlers.

hawkinsp committed 290 days ago

0f2b3f21

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

Jake VanderPlas committed 290 days ago

c6a8050c

Older