ROCm/jax

Pull Requests Commits

adding jax unit test tridiagonal for rocm platform

zahiqbal committed 249 days ago

ba2a99df

Fix test flakyness by blocking until the data is ready.

pschuh committed 250 days ago

55200553

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

chr1sj0nes committed 250 days ago

2beff6a1

remove inaccurate inline comment in `PRNGKeyArray` constructor

froystig committed 250 days ago

90af5977

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

froystig committed 250 days ago

47bc2f55

Merge pull request #27984 from carlosgmartin:logsumexp_doc

Google-ML-Automation committed 250 days ago

25e0fe59

Merge pull request #28047 from jakevdp:logsoftmax-dep

Google-ML-Automation committed 250 days ago

002be7a1

Enable standard_insert_pvary for optimization_barrier which was disabled before.

yashk2810 committed 250 days ago

655bfcac

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

bythew3i committed 250 days ago

73889130

Clean up softmax initial deprecation

jakevdp committed 250 days ago

b271a67b

Fix bugs in tp_traverse handlers.

hawkinsp committed 250 days ago

393c555a

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

Jake VanderPlas committed 250 days ago

ba887778

[NFC] Rename `standard_insert_pbroadcast` to `standard_insert_pvary`

yashk2810 committed 250 days ago

6e00b5e0

[Mosaic:TPU] Fix bug in `rotateVregRows`

tlongeri committed 250 days ago

c527ddb7

Merge pull request #28022 from jakevdp:bernoulli-fix

Google-ML-Automation committed 250 days ago

b336daf7

Merge pull request #28036 from ROCm:ci_fix_rnn_lowering-upstream

Google-ML-Automation committed 250 days ago

87e4b5f8

Fix lowering code for ROCm RNN

Ruturaj4 committed 250 days ago

34c2dbfd

Merge pull request #28027 from jax-ml:explicit_axes

Google-ML-Automation committed 250 days ago

3b359ba2

jax.random.bernoulli: use higher-resolution sampler

jakevdp committed 250 days ago

c56cf4f6

Update XLA dependency to use revision

Google-ML-Automation committed 250 days ago

4d692d15

Deprecate the remaining exports from jax.lib.xla_client.

hawkinsp committed 250 days ago

0b047396

[CI] Propagate halt connection to tpu tests

MichaelHudgins committed 250 days ago

06a77b7b

[CI] Update GPU optional presubmit naming

MichaelHudgins committed 251 days ago

aed3297f

Add explicit_axes section to the doc

yashk2810 committed 251 days ago

09edc494

[pallas] Fix spelling of 'fusible'.

chr1sj0nes committed 251 days ago

1926b99b

Adds a debugging message to assert, otherwise the error is pretty cryptic.

marksandler2 committed 251 days ago

0ed0fb7c

[Pallas/Fuser] Add basic closed over consts support to pull_block_spec

sharadmv committed 251 days ago

4fa3cd91

Merge pull request #27999 from vfdev-5:fix-skip-test-pattern-with-opt-deps

Google-ML-Automation committed 251 days ago

69d21c69

Deprecate the contents of jax.util.

hawkinsp committed 251 days ago

57e33bcb

Fixed the way to skip tests using optional python dependencies

vfdev-5 committed 251 days ago

7b4b2f47

Older