ROCm/jax

Pull Requests Commits

[ROCm]: Add support to continue on fail

Rahul Batra committed 2 years ago

7559b07a

Use opaque layout PJRT_Layouts_MemoryLayout in PjRtCApiBuffer::layout() to keep all the layout information.

Jieying Luo committed 2 years ago

51e59518

Move throwing the error for invalid compute_on into the try-finally context.

pschuh committed 2 years ago

b47157db

Remove warning logs for primary_host/remote_storage "incompatibility".

cpgaffney1 committed 2 years ago

26b4848f

Merge pull request #21467 from jakevdp:jnp-bincount-doc

a-googler committed 2 years ago

8c55e8bc

creat jax config api to allow custom pjrt client create option settings. this allows a device platform's pjrt client be aware of the calling (customer) ml framework

zuasia committed 2 years ago

91d68b55

Improve documentation for jnp.bincount

jakevdp committed 2 years ago

fcdbb2a2

Merge pull request #21466 from jakevdp:warning-filters

a-googler committed 2 years ago

db118423

Merge pull request #21457 from jakevdp:ravel-index-doc

a-googler committed 2 years ago

4cb6a150

CI: remove stale warning filters

jakevdp committed 2 years ago

b441a09a

Merge pull request #21463 from jakevdp:jnp-take-doc

a-googler committed 2 years ago

4461b7fe

Merge pull request #20705 from chaserileyroberts:chase/pbroadcast_channel_fix

a-googler committed 2 years ago

a7bce471

Improve documentation for jnp.take & jnp.take_along_axis

jakevdp committed 2 years ago

9e80c307

Merge pull request #21462 from sritchie:sritchie-patch-1

a-googler committed 2 years ago

2b6bcb55

Fix pip dependency after cuDNN 9 upgrade

beckerhe committed 2 years ago

43726365

Clean up version switches from dense array migration

Michael Levesque-Dion committed 2 years ago

43f51d73

[Mosaic] Add relayout for (1, 128 * packing) -> (packing, 128).

tlongeri committed 2 years ago

8b958536

Update quickstart to suggested jax[cuda12] installation

sritchie committed 2 years ago

Verified 822bc364

Pipe channel handle

chaserileyroberts committed 2 years ago

af6970e4

Merge pull request #21452 from jakevdp:take-fill-value

a-googler committed 2 years ago

8fec5d6c

Merge pull request #21193 from sh0416:main

a-googler committed 2 years ago

82ad1daa

Merge pull request #21453 from jakevdp:index-update-annotation

a-googler committed 2 years ago

83871d3b

Improve docstrings for unravel_index & ravel_multi_index

jakevdp committed 2 years ago

6803d771

Update the deprecation message of `backend` and `device` argument of `jit` to be more actionable.

yashk2810 committed 2 years ago

ff3db9b3

Merge pull request #21451 from gnecula:poly_pmap

a-googler committed 2 years ago

de28ee6f

Fix type annotation for index-update

jakevdp committed 2 years ago

3ade0232

jnp.take_along_axis: support fill_value

jakevdp committed 2 years ago

2d23a66c

Merge pull request #21376 from ROCm:ci_f8

a-googler committed 2 years ago

720d2b87

Merge pull request #21434 from jakevdp:take-annotation

a-googler committed 2 years ago

ebd9c740

[export] Fix calling under pmap of exported computation with polymorphic shapes

gnecula committed 2 years ago

acb56a29

Older