ROCm/jax

Pull Requests Commits

Update rocm-test-requirements.txt

gulsumgudukbay committed 319 days ago

9e5c47cb

Revert "fixed typo"

gulsumgudukbay committed 319 days ago

69ec57e6

gulsumgudukbay committed 319 days ago

d2b14cb4

[0.6.0-UT] adding a QA test requirements file for the run_single_gpu test script… (#446)

gulsumgudukbay committed 320 days ago

Verified 1fc5d2a5

[0.6.0-UT] Skip test_dot for ROCm <6.5, and shapes exceed shared memory limits (#448)

psanal35 committed 320 days ago

Verified d6456448

[0.6.0-UT] Add shared utility function get_rocm_version to test_util.py

psanal35 committed 320 days ago

Verified 893c907c

unskip test_select_n as it is passing on 0.6.0

gulsumgudukbay committed 324 days ago

5b38406a

added conditional seq_len assignment for lower dimension for tests/pallas/gpu_ops_test.py::FusedAttentionTest::test_fused_attention_[fwd,bwd] tests with the rationale. The test passes now because the input fits in shared memory. before, the tests were failing due to the seq exceeding the shared memory size and causing XLARuntimeError.

gulsumgudukbay committed 325 days ago

5659c086

[JAX] Fix numerical instability in testEighTinyNorm for ROCm GPU

phambinhfin committed 325 days ago

3c6ab9f3

JAX release v0.6.0

hawkinsp committed 353 days ago

9e9a3f03

[Pallas] Fix relative links in pipelining docs

justinjfu committed 353 days ago

28538e4c

Fix bug in argnums_partial_except when static_argnums is unsorted.

dfm committed 353 days ago

34234010

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

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

33d43f7a

[Pallas] Move pipelining doc figures to static folder

justinjfu committed 353 days ago

a9465e70

Add explicit sharding to Introduction to parallel programming doc

yashk2810 committed 353 days ago

24fdc0ad

Update XLA dependency to use revision

Google-ML-Automation committed 353 days ago

6b94a5ce

[Pallas] Refactor pipelining docs.

justinjfu committed 353 days ago

17e2dcaf

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

nitins17 committed 353 days ago

bc4fcd87

disable packed layout test on old arch prior to Hopper

Cjkkkk committed 353 days ago

fb875101

Update XLA dependency to use revision

Google-ML-Automation committed 353 days ago

a5236392

Fix race suppression for 3.13

vfdev-5 committed 353 days ago

8d474a3f

Add basic documentation for Pallas:MGPU layouts and transforms

apaszke committed 353 days ago

273e6efc

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

superbobry committed 353 days ago

bb6ca370

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

apaszke committed 353 days ago

b7a041cc

Add docs for GPU memory spaces in Pallas:MGPU

apaszke committed 353 days ago

fe96e029

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

bythew3i committed 353 days ago

3a54f93d

Fix test flakyness by blocking until the data is ready.

pschuh committed 353 days ago

3ad33334

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

chr1sj0nes committed 353 days ago

3dfbdd09

remove inaccurate inline comment in `PRNGKeyArray` constructor

froystig committed 353 days ago

be3c4d70

Newer Older