ROCm/jax

Pull Requests Commits

Remove ROCm 6.2 build

charleshofer committed 253 days ago

4ac6032b

Account for versioned clang binaries

charleshofer committed 253 days ago

fb0bc234

JAX release v0.6.0

hawkinsp committed 258 days ago

9e9a3f03

[Pallas] Fix relative links in pipelining docs

justinjfu committed 258 days ago

28538e4c

Fix bug in argnums_partial_except when static_argnums is unsorted.

dfm committed 258 days ago

34234010

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

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

33d43f7a

[Pallas] Move pipelining doc figures to static folder

justinjfu committed 258 days ago

a9465e70

Add explicit sharding to Introduction to parallel programming doc

yashk2810 committed 258 days ago

24fdc0ad

Update XLA dependency to use revision

Google-ML-Automation committed 258 days ago

6b94a5ce

[Pallas] Refactor pipelining docs.

justinjfu committed 258 days ago

17e2dcaf

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

nitins17 committed 258 days ago

bc4fcd87

disable packed layout test on old arch prior to Hopper

Cjkkkk committed 258 days ago

fb875101

Update XLA dependency to use revision

Google-ML-Automation committed 258 days ago

a5236392

Fix race suppression for 3.13

vfdev-5 committed 258 days ago

8d474a3f

Add basic documentation for Pallas:MGPU layouts and transforms

apaszke committed 258 days ago

273e6efc

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

superbobry committed 258 days ago

bb6ca370

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

apaszke committed 258 days ago

b7a041cc

Add docs for GPU memory spaces in Pallas:MGPU

apaszke committed 258 days ago

fe96e029

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

bythew3i committed 258 days ago

3a54f93d

Fix test flakyness by blocking until the data is ready.

pschuh committed 258 days ago

3ad33334

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

chr1sj0nes committed 258 days ago

3dfbdd09

remove inaccurate inline comment in `PRNGKeyArray` constructor

froystig committed 258 days ago

be3c4d70

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

froystig committed 258 days ago

b35f5f7d

Minor improvements to doc for jax.nn.logsumexp.

carlosgmartin committed 258 days ago

b13364fb

Clean up softmax initial deprecation

jakevdp committed 258 days ago

1aeac03c

Enable standard_insert_pvary for optimization_barrier which was disabled before.

yashk2810 committed 258 days ago

c3d8ac79

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

bythew3i committed 258 days ago

5214342e

Fix bugs in tp_traverse handlers.

hawkinsp committed 258 days ago

0f2b3f21

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

Jake VanderPlas committed 258 days ago

c6a8050c

Older