ROCm/jax

Pull Requests Commits

[internal] add deprecation test utilities

jakevdp committed 1 year ago

1ee01567

Merge pull request #25741 from jakevdp:solve-dep

Google-ML-Automation committed 1 year ago

5d0ee432

Merge pull request #25791 from mattjj:logsumexp-where-grad-nan

Google-ML-Automation committed 1 year ago

aed79707

Annotate several tests as thread-unsafe.

hawkinsp committed 1 year ago

8f2f4b45

Merge pull request #25787 from dfm:tri-diag-jvp

Google-ML-Automation committed 1 year ago

016fca79

Merge pull request #25838 from jakevdp:simplify-nightly-reporting

Google-ML-Automation committed 1 year ago

f7b6deeb

jnp.linalg.solve: finalize deprecation of batched 1D solves

jakevdp committed 1 year ago

051abafd

Update XLA dependency to use revision

Google-ML-Automation committed 1 year ago

1cc07dd3

Simplify failure reporting for nightly CI job

jakevdp committed 1 year ago

7f55ccca

Add a JVP rule for lax.linalg.tridiagonal_solve + various fixes.

dfm committed 1 year ago

167a48f6

Activate FFI implementation of tridiagonal reduction on GPU.

dfm committed 1 year ago

39ce7916

Add LAPACK lowering for lax.linalg.tridiagonal_solve on CPU.

dfm committed 1 year ago

c1de7c73

Merge pull request #25831 from jax-ml:avoid-float0-tracers

Google-ML-Automation committed 1 year ago

4f106b8a

[Mosaic GPU] Add initial layout mismatch resolution for `splat`/`strided` layouts.

bchetioui committed 1 year ago

1893881b

Merge pull request #20282 from tttc3:pivoted-qr

Google-ML-Automation committed 1 year ago

564b6b0d

Avoid creating float0 JVPTracers

dougalm committed 1 year ago

ba9b2ca5

Merge pull request #25771 from dfm:custom-partition-ffi-tutorial

Google-ML-Automation committed 1 year ago

1fe72ee8

[Mosaic TPU] Add support for integer truncation from packed types

apaszke committed 1 year ago

d2a5e8d0

Make JAX test suite pass (at least most of the time) with multiple threads enabled.

hawkinsp committed 1 year ago

c61b2f6b

Skip RnnTest.test_struct_encoding_determinism.

belitskiy committed 1 year ago

86643a1b

[pallas:mosaic_gpu] Tests now pass with x64 enabled

superbobry committed 1 year ago

18018d9c

[Pallas] Improve testing for lowering of dtype conversions + fix uncovered bugs

apaszke committed 1 year ago

74cf67df

Reverts dbe9ccd6dccd83c365021677c7e17e843d4559c4

chr1sj0nes committed 1 year ago

a27566cc

Merge pull request #25800 from gnecula:improve_error_switch

Google-ML-Automation committed 1 year ago

8c236898

Merge pull request #25797 from gnecula:print_debug_info

Google-ML-Automation committed 1 year ago

228d3cef

[better_errors] Improve error message for lax.switch branches output structure mismatch

gnecula committed 1 year ago

c2adfbf1

[aot] Add support for as_text(debug_info=True).

gnecula committed 1 year ago

dd0447a7

Remove extraneous print statement in a test

yashk2810 committed 1 year ago

6319126e

Merge pull request #25803 from sergachev:fix_rnn_desc

Google-ML-Automation committed 1 year ago

061408ac

#sdy enable pure callbacks and debug prints in JAX.

bartchr808 committed 1 year ago

dc53c563

Newer Older