ROCm/jax

Pull Requests Commits

Enable remote compilation for jax

alekstheod committed 3 days ago

601a9a91

Fix precision of jax.scipy.special.ndtr in the negative tail.

hawkinsp committed 3 days ago

03930d1f

Reverts 5ab78d311eb8fd5bd1512a51bcf2c8c766df09f7

jerryx99 committed 3 days ago

ff992d6f

Merge pull request #41938 from mattjj:push-ulnpouqqoouq

Google-ML-Automation committed 3 days ago

7231eef5

[nn] Mask tangents and cotangents with where in _softmax_jvp

Jake VanderPlas CA committed 3 days ago

b9ee546e

Switch from libstdc++ to statically linked libc++

yuriivcs committed 3 days ago

41b77aff

[pallas:fuser] Check pulled kernel-output block specs in push_pull_block_spec

rdyro committed 3 days ago

65054360

Merge pull request #42086 from jakevdp:nn-test-softmax

Google-ML-Automation committed 3 days ago

2266f394

Lower scan_p with unreduced scanned arrays via HiJAX.

yashk2810 committed 3 days ago

1531d80f

[test] fix test broken under softmax_custom_jvp

jakevdp committed 3 days ago

217f98ab

[pallas:tpu] Document `sem=None`, `src_sem`, `wait_read()` and `wait_write()` on SparseCore

rdyro committed 3 days ago

9b2ff832

[jax-buildcop] Skip `lobpcg_test` on CPU under ThreadSanitizer where it times out

Google-ML-Automation committed 3 days ago

70fd33dc

Merge pull request #42070 from jakevdp:softmax-initial

Google-ML-Automation committed 3 days ago

28fb3ea0

[Mosaic] Remove JAX Mosaic TPU forwarding files and build targets

WillFroom committed 3 days ago

2e61e964

Lower unreduced reduce_sum, reduce_max, reduce_min, and copy via HiJAX.

yashk2810 committed 3 days ago

6ea19e0e

[jax-buildcop] Disable lax_numpy_indexing_test on CPU under ASAN.

Google-ML-Automation committed 3 days ago

2da81a55

[jax-buildcop] Seed RNG and pass SVD tolerances to _CompileAndCheck in SvdTest

Google-ML-Automation committed 3 days ago

9cf555f3

[CI] Update CI zizmor to v1.31.1; use Zizmor action (v0.8.0) for auto-update PRs.

belitskiy committed 3 days ago

4c6780a4

Lower standard lax, linear, and shape primitives for HiJAX unreduced.

yashk2810 committed 3 days ago

7983b975

[Mosaic GPU] Support `.sync_restrict::shared::read::mma::a` in `tcgen05.commit_arrive` on `sm_107f`.

bchetioui committed 3 days ago

b8a8d87a

[JAX] Add F8E4M3FN_F8E4M3FN_F32_X{3,4} DotAlgorithmPreset values

SamGinzburg committed 3 days ago

b2de65e3

Remove unused argument in internal softmax

jakevdp committed 3 days ago

e1cf13e7

[Mosaic GPU] Validate and lower TMA async copies in transformed space.

golechwierowicz committed 3 days ago

ec456bfe

[Pallas:MGPU] Derive `a_aval.dtype` and `b_aval.dtype` from `_handle_transforms` in `_tcgen05_mma_lowering`.

allanrenucci committed 3 days ago

3738fa5c

[Mosaic GPU] Support non-power-of-two vector bitwidths in `getelementptr` and loads/stores.

golechwierowicz committed 3 days ago

65bf1782

[JAX] Use WEAKORDER comparison type for floating-point sort comparators

akuegel committed 3 days ago

5ab78d31

Update XLA dependency to use revision http://github.com/openxla/xla/commit/d5229a92b1d15d30ba87f68c42e44cae0e77f54c

Google-ML-Automation committed 3 days ago

3e6335a8

Don't allow arbitrary unreduced reshard in convert_element_type. The current rule is similar to broadcast_in_dim's rule.

yashk2810 committed 3 days ago

70b48fe1

Add GPU kernels for lgamma and digamma.

hawkinsp committed 4 days ago

745e9e3c

Fix some unreduced tests when hijax_unreduced=True

yashk2810 committed 4 days ago

7780bb2f

Older