Skip to content

Pull requests: jax-ml/jax

New pull request New
405 Open 15,890 Closed
405 Open 15,890 Closed
Author
Filter by author
Loading
Label
Filter by label
Loading
Use alt + click/return to exclude labels
or ⇧ + click/return for logical OR
Projects
Filter by project
Loading
Milestones
Filter by milestone
Loading
Reviews
Filter by reviews
No reviews Review required Approved review Changes requested
Assignee
Filter by who’s assigned
Sort
Sort by
Newest Oldest Most commented Least commented Recently updated Least recently updated Best match
Most reactions
👍 👎 😄 🎉 😕 ❤️ 🚀 👀

Pull requests list

Update jax.grad docs
#24924 opened Nov 15, 2024 by 8bitmp3 • Draft
Update jax.random docs
#24923 opened Nov 15, 2024 by 8bitmp3 • Draft
[sharding_in_types] Don't emit a wsc under full manual mode to avoid increasing HLO size by a lot
#24922 opened Nov 15, 2024 by copybara-service bot Loading…
@yashk2810
Update jax.jit docs
#24921 opened Nov 15, 2024 by 8bitmp3 • Draft
5
Update jax.numpy docs
#24920 opened Nov 15, 2024 by 8bitmp3 • Draft
[Mosaic TPU] Fold sublane offset to indices when storing to untiled ref.
#24919 opened Nov 15, 2024 by copybara-service bot Loading…
Consolidate material on PRNGs and add a short summary to Key Concepts. pull ready Ready for copybara import and testing
#24917 opened Nov 15, 2024 by emilyfertig Loading…
@jakevdp
8
DOC: update main landing page style
#24916 opened Nov 15, 2024 by jakevdp Loading…
@jakevdp
Set __module__ attribute for objects in jax.numpy
#24912 opened Nov 15, 2024 by jakevdp Loading…
@jakevdp
[AutoPGLE] Temporary disable pgle_test in the OSS.
#24911 opened Nov 15, 2024 by copybara-service bot Loading…
Add JAX_COMPILATION_CACHE_EXPECT_PGLE option
#24910 opened Nov 15, 2024 by olupton Loading…
C++ tree with path API
#24906 opened Nov 14, 2024 by copybara-service bot Loading…
@IvyZX
Make logaddexp and logaddexp2 into ufuncs pull ready Ready for copybara import and testing
#24903 opened Nov 14, 2024 by jakevdp Loading…
@jakevdp
2
[Pallas TPU] Support masked store
#24902 opened Nov 14, 2024 by copybara-service bot Loading…
[pallas] Minor simplifications to Pallas interpreter.
#24900 opened Nov 14, 2024 by copybara-service bot Loading…
@chr1sj0nes
Added CI job with TSAN and free-threading
#24898 opened Nov 14, 2024 by vfdev-5 • Draft
[Mosaic TPU] Add general tpu.vector_store and support masked store.
#24892 opened Nov 14, 2024 by copybara-service bot Loading…
[Mosaic TPU] Support relayout for mask vector
#24891 opened Nov 13, 2024 by copybara-service bot Loading…
Implement lax.pad in Pallas.
#24890 opened Nov 13, 2024 by copybara-service bot Loading…
[Pallas] Increase test coverage of pl.dot.
#24880 opened Nov 13, 2024 by copybara-service bot Loading…
[JAX/version] Support adding custom version suffix.
#24861 opened Nov 12, 2024 by yliu120 Loading…
2
[JAX/distributed] Small clean ups on warning emissions in jax.distributed.
#24853 opened Nov 12, 2024 by yliu120 Loading…
Bump hypothesis from 6.102.4 to 6.118.7 dependencies Pull requests that update a dependency file python Pull requests that update Python code
#24840 opened Nov 11, 2024 by dependabot bot Loading…
Improvements in the handling of tracing debug info pull ready Ready for copybara import and testing
#24831 opened Nov 11, 2024 by gnecula • Draft
@gnecula
Pattern match dot algorithm spec to preset name pull ready Ready for copybara import and testing
#24820 opened Nov 9, 2024 by dfm Loading…
@dfm
ProTip! Mix and match filters to narrow down what you’re looking for.