ROCm/jax

Pull Requests Commits

adding jax unit test tridiagonal for rocm platform

zahiqbal committed 1 year ago

ba2a99df

Fix test flakyness by blocking until the data is ready.

pschuh committed 1 year ago

55200553

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

chr1sj0nes committed 1 year ago

2beff6a1

remove inaccurate inline comment in `PRNGKeyArray` constructor

froystig committed 1 year ago

90af5977

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

froystig committed 1 year ago

47bc2f55

Merge pull request #27984 from carlosgmartin:logsumexp_doc

Google-ML-Automation committed 1 year ago

25e0fe59

Merge pull request #28047 from jakevdp:logsoftmax-dep

Google-ML-Automation committed 1 year ago

002be7a1

Enable standard_insert_pvary for optimization_barrier which was disabled before.

yashk2810 committed 1 year ago

655bfcac

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

bythew3i committed 1 year ago

73889130

Clean up softmax initial deprecation

jakevdp committed 1 year ago

b271a67b

Fix bugs in tp_traverse handlers.

hawkinsp committed 1 year ago

393c555a

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

Jake VanderPlas committed 1 year ago

ba887778

[NFC] Rename `standard_insert_pbroadcast` to `standard_insert_pvary`

yashk2810 committed 1 year ago

6e00b5e0

[Mosaic:TPU] Fix bug in `rotateVregRows`

tlongeri committed 1 year ago

c527ddb7

Merge pull request #28022 from jakevdp:bernoulli-fix

Google-ML-Automation committed 1 year ago

b336daf7

Merge pull request #28036 from ROCm:ci_fix_rnn_lowering-upstream

Google-ML-Automation committed 1 year ago

87e4b5f8

Fix lowering code for ROCm RNN

Ruturaj4 committed 1 year ago

34c2dbfd

Merge pull request #28027 from jax-ml:explicit_axes

Google-ML-Automation committed 1 year ago

3b359ba2

jax.random.bernoulli: use higher-resolution sampler

jakevdp committed 1 year ago

c56cf4f6

Update XLA dependency to use revision

Google-ML-Automation committed 1 year ago

4d692d15

Deprecate the remaining exports from jax.lib.xla_client.

hawkinsp committed 1 year ago

0b047396

[CI] Propagate halt connection to tpu tests

MichaelHudgins committed 1 year ago

06a77b7b

[CI] Update GPU optional presubmit naming

MichaelHudgins committed 1 year ago

aed3297f

Add explicit_axes section to the doc

yashk2810 committed 1 year ago

09edc494

[pallas] Fix spelling of 'fusible'.

chr1sj0nes committed 1 year ago

1926b99b

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

marksandler2 committed 1 year ago

0ed0fb7c

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

sharadmv committed 1 year ago

4fa3cd91

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

Google-ML-Automation committed 1 year ago

69d21c69

Deprecate the contents of jax.util.

hawkinsp committed 1 year ago

57e33bcb

Fixed the way to skip tests using optional python dependencies

vfdev-5 committed 1 year ago

7b4b2f47

Older