ROCm/jax

Pull Requests Commits

Merge pull request #33636 from jakevdp:doc-extend-backend

Google-ML-Automation committed 93 days ago

cb2c161d

Merge pull request #33667 from jakevdp:checkpoint-name-doc

Google-ML-Automation committed 93 days ago

ff282428

[doc] add API docs for jax.extend.backend

jakevdp committed 93 days ago

f1c66433

Add support for NumPy 2.4 __numpy_dtype__ protocol

jakevdp committed 93 days ago

05d50a36

[doc] better docstring for checkpoint_name()

jakevdp committed 93 days ago

f934d7e6

jax.checkpoint: deprecate the concrete argument

jakevdp committed 93 days ago

de776ac3

Merge pull request #33606 from partev:patch-1

Google-ML-Automation committed 93 days ago

176e3cb9

[Take 2] Merge all_gather with all_gather_reduced, psum_scatter with unreduced_psum_scatter and psum with unreduced_psum.

yashk2810 committed 93 days ago

ad44116f

Reverts 39f36f535ec20e01a2d6a9323816fd8758bca38b

Google-ML-Automation committed 93 days ago

5b9cfa39

Merge pull request #33661 from jax-ml:more-traceback-regression-tests

Google-ML-Automation committed 93 days ago

e7046392

[no-thunks] Add more traceback regression tests

dougalm committed 93 days ago

4871fcaa

Merge pull request #33660 from jakevdp:ad_checkpoint

Google-ML-Automation committed 93 days ago

e4cadda6

Merge all_gather with all_gather_reduced, psum_scatter with unreduced_psum_scatter and psum with unreduced_psum.

yashk2810 committed 93 days ago

39f36f53

Deprecate jax.ad_checkpoint.remat and jax.ad_checkpoint.checkpoint

jakevdp committed 93 days ago

360a571a

Merge pull request #33634 from jakevdp:ad-checkpoint-doc

Google-ML-Automation committed 93 days ago

7a395f31

If LAPACK kernels are linked statically, don't override them from scipy.

hawkinsp committed 93 days ago

c1adda87

Merge pull request #33663 from jakevdp:checkpoint-doc

Google-ML-Automation committed 93 days ago

a915a311

Merge pull request #33551 from ashahba:ashahba/fix-cpu-noteook

Google-ML-Automation committed 93 days ago

a3069ac9

[doc] add JAX version context to JEP 11830

jakevdp committed 93 days ago

1a75f103

Make every parallel primitive multiple_results=False. Primitives include `psum_p`, `psum_invariant_p`, `pmax_p`, `pmin_p`, `pvary_p`, `preduced_p`, `unreduced_psum_p`, `vary_unreduced_cast_p` and `reduced_vary_cast_p`

yashk2810 committed 93 days ago

f43e21c6

[doc] Add missing documentation for jax.ad_checkpoint

jakevdp committed 93 days ago

3f0d0f8e

Merge pull request #33640 from jax-ml:scan-stack-frames

Google-ML-Automation committed 94 days ago

ad292f0d

[no-thunks] Reduce stack frames when tracing initial-style primitives.

dougalm committed 94 days ago

b2c53a3e

[Mosaic GPU] Impelment basic support for `memref.expand_shape` in WG semantics.

dimitar-asenov committed 94 days ago

97443af2

Remove an apparently stale test skip.

hawkinsp committed 94 days ago

d0e62bea

Merge pull request #33564 from ROCm:ci_add_pallas_triton_isfinite

Google-ML-Automation committed 94 days ago

5a2f70a5

[pallas:sc] Bumped the number of shards for :tpu_sparsecore_pallas_test

superbobry committed 94 days ago

8f8daebb

Jax changes before updating the stack frame format for XLA.

pifon2a committed 94 days ago

b432bb15

[Pallas:SC] Fix a bug where get_global lowering depended indirectly on sortable dict keys (via tree.flatten) by making MemoryRef comparable.

brianwa84 committed 94 days ago

8e8c79a4

[Pallas:SC] Add support for neg_p, abs_p.

brianwa84 committed 94 days ago

1f1f1e9c

Newer Older