ROCm/jax

Pull Requests Commits

Use hermetic Python in JAX, see "Managing hermetic Python" in developer.md for details

vam-google committed 1 year ago

517e299a

[Mosaic GPU] Adjust memref.expand_shape construction to pass in the new args

Google-ML-Automation committed 1 year ago

a820387a

Merge pull request #21251 from jakevdp:fix-pmap-warning

a-googler committed 1 year ago

6a949ae9

Fix warning filter in test_pmap_of_prng_key

jakevdp committed 1 year ago

f5a59ccd

Merge pull request #21247 from dfm:gh20481

a-googler committed 1 year ago

15381972

[Mosaic GPU] Fix argument checking in case the input shape is not a tuple.

Google-ML-Automation committed 1 year ago

bb400756

Add informative error for invalid unroll in scan

dfm committed 1 year ago

09a4b38a

Set V3 cores to 4

MichaelHudgins committed 1 year ago

Verified 76a7b19c

Run v3-8 tests with cores set at 8

MichaelHudgins committed 1 year ago

Verified 0232cb9f

Add missing bash symbol

MichaelHudgins committed 1 year ago

Verified 181da128

Add core count to tpu nightly fix v5 job

MichaelHudgins committed 1 year ago

Verified 30156999

Merge pull request #21237 from gnecula:clean_tokens

a-googler committed 1 year ago

4ccac4c6

[Mosaic GPU] Explicitly check the smem size so that we get a good error rather than a cryptic cuda error

Google-ML-Automation committed 1 year ago

11a45138

[Mosaic GPU] Fix np.random API mistake in tests

Google-ML-Automation committed 1 year ago

874a20a2

Fix CI test failure in line_search_test.

hawkinsp committed 1 year ago

7653db80

Cleanup token handling during lowering

gnecula committed 1 year ago

41153b16

Reverts 9e7830df2df9362edcf2e18e353d327fdecae678

yueshengys committed 1 year ago

66a92c41

Update XLA dependency to use revision

a-googler committed 1 year ago

b8cec530

Add doc for jax.experimental.serialize_executable module

nouiz committed 1 year ago

936ec983

Merge pull request #21222 from tempoxylophone:patch-1

a-googler committed 1 year ago

9af0a317

Merge pull request #21221 from nouiz:doc

a-googler committed 1 year ago

48b90181

Merge pull request #21226 from Micky774:linalg-depr

a-googler committed 1 year ago

1dff5b37

Merge pull request #21227 from jakevdp:eye-offset

a-googler committed 1 year ago

d4f36e25

Cholesky rank-1 update kernel for JAX.

a-googler committed 1 year ago

e8b06ccf

Added a very rough sketch of Mosaic GPU lowering for Pallas

superbobry committed 1 year ago

e2918ca1

jnp.eye: allow k to be dynamic

jakevdp committed 1 year ago

3024c782

Add support for i1 vmasks with packed tiling and 16-bit comparisons (requires hardware support)

tlongeri committed 1 year ago

0ad5167d

Rename rcond/tol to rtol in linalg.matrix_rank and linalg.pinv

Micky774 committed 1 year ago

5cc255b7

Merge pull request #21223 from jakevdp:kl-div-doc

a-googler committed 1 year ago

0501d3d7

Finalize deprecations of several APIs

Jake VanderPlas committed 1 year ago

bb5787da

Newer Older