ROCm/jax

Pull Requests Commits

Revised supported gfx-arch, fixed build script

zahiqbal committed 1 year ago

88c9b5dc

Fixes for ci build

mrodden committed 1 year ago

a7833b89

JAX release v0.6.0

hawkinsp committed 1 year ago

9e9a3f03

[Pallas] Fix relative links in pipelining docs

justinjfu committed 1 year ago

28538e4c

Fix bug in argnums_partial_except when static_argnums is unsorted.

dfm committed 1 year ago

34234010

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

yashk2810 committed 1 year 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 1 year ago

33d43f7a

[Pallas] Move pipelining doc figures to static folder

justinjfu committed 1 year ago

a9465e70

Add explicit sharding to Introduction to parallel programming doc

yashk2810 committed 1 year ago

24fdc0ad

Update XLA dependency to use revision

Google-ML-Automation committed 1 year ago

6b94a5ce

[Pallas] Refactor pipelining docs.

justinjfu committed 1 year ago

17e2dcaf

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

nitins17 committed 1 year ago

bc4fcd87

disable packed layout test on old arch prior to Hopper

Cjkkkk committed 1 year ago

fb875101

Update XLA dependency to use revision

Google-ML-Automation committed 1 year ago

a5236392

Fix race suppression for 3.13

vfdev-5 committed 1 year ago

8d474a3f

Add basic documentation for Pallas:MGPU layouts and transforms

apaszke committed 1 year ago

273e6efc

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

superbobry committed 1 year ago

bb6ca370

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

apaszke committed 1 year ago

b7a041cc

Add docs for GPU memory spaces in Pallas:MGPU

apaszke committed 1 year ago

fe96e029

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

bythew3i committed 1 year ago

3a54f93d

Fix test flakyness by blocking until the data is ready.

pschuh committed 1 year ago

3ad33334

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

chr1sj0nes committed 1 year ago

3dfbdd09

remove inaccurate inline comment in `PRNGKeyArray` constructor

froystig committed 1 year ago

be3c4d70

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

froystig committed 1 year ago

b35f5f7d

Minor improvements to doc for jax.nn.logsumexp.

carlosgmartin committed 1 year ago

b13364fb

Clean up softmax initial deprecation

jakevdp committed 1 year ago

1aeac03c

Enable standard_insert_pvary for optimization_barrier which was disabled before.

yashk2810 committed 1 year ago

c3d8ac79

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

bythew3i committed 1 year ago

5214342e

Fix bugs in tp_traverse handlers.

hawkinsp committed 1 year ago

0f2b3f21

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

Jake VanderPlas committed 1 year ago

c6a8050c

Older