ROCm/jax

Pull Requests Commits

Add defaults for ROCm build vars

mrodden committed 1 year ago

d7481d2d

Add copyright and license headers to new files

mrodden committed 1 year ago

04373f20

Add support for ROCm development builds

mrodden committed 1 year ago

b92f5ec9

Add internal release support to get_rocm.py

mrodden committed 1 year ago

f8362e38

[ROCm] Change ROCm builds to manylinux wheels

mrodden committed 1 year ago

3a1601dd

Fix lint in run_single_gpu.py

yueshengys committed 1 year ago

91079db3

[ROCM] add pytest reportlog

Ruturaj4 committed 1 year ago

534b7a33

[JAX] Fix run_single_gpu.py to include correct paths

Ruturaj4 committed 1 year ago

95ae95ca

[ROCm]: Add support to continue on fail

Rahul Batra committed 1 year ago

84a3e91c

Start JAX and jaxlib 0.4.31 release

yashk2810 committed 1 year ago

7fd93027

Merge pull request #22703 from Rifur13:plugin-fix

a-googler committed 1 year ago

f070c065

Merge pull request #19760 from Blair-Johnson:fix-pytree-grads-sparse

a-googler committed 1 year ago

9beb4f14

Re-pack gradients of jax.experimental.sparse.grad() to match original pytrees & test cases

Blair-Johnson committed 1 year ago

802a14cd

Merge pull request #22690 from jakevdp:inplace-doc

a-googler committed 1 year ago

85e83b50

Bump the minimum CUDNN version to v9.1.

hawkinsp committed 1 year ago

d1c0d993

Ignore the Deprecation warning produced about `native_serialization=False`.

belitskiy committed 1 year ago

6127baa1

Bump minimum SciPy version to 1.10.

hawkinsp committed 1 year ago

fd23b873

Improve documentation for jnp.put, jnp.place, jnp.fill_diagonal

jakevdp committed 1 year ago

cfa1e785

Skip cuda backend initialization if no nvidia GPUs are visible.

Rifur13 committed 1 year ago

0224235e

Merge pull request #22593 from gnecula:pallas_more_simplification

a-googler committed 1 year ago

e78e643b

Rearrange the LAPACK handler definitions in jaxlib to avoid duplicate handler errors.

dfm committed 1 year ago

ff4e0b12

Skip tests/mock_gpu_test.py on pytest.

belitskiy committed 1 year ago

fef91fb2

[pallas] More simplification of grid mapping and calling convention

gnecula committed 1 year ago

70a11acb

[pallas] Add lowering errors for block shapes that are not supported.

gnecula committed 1 year ago

68972de0

Reduced the input size in PallasCallInputOutputAliasingTest

superbobry committed 1 year ago

ccc4c42e

[Mosaic GPU] Fix two subtle issues with kernel lowering

apaszke committed 1 year ago

a00b659b

Update XLA dependency to use revision

Google-ML-Automation committed 1 year ago

6a7822a7

Update XLA dependency to use revision

Google-ML-Automation committed 1 year ago

74649be7

Finalize deprecation of jax.random.shuffle

Jake VanderPlas committed 1 year ago

a17c8d94

Merge pull request #22684 from froystig:rngdoc

a-googler committed 1 year ago

dab15d6f

Newer Older