Skip to content

Pull requests: google/jax

New pull request New
307 Open 13,195 Closed
307 Open 13,195 Closed
Author
Filter by author
Label
Filter by label
Use alt + click/return to exclude labels
or ⇧ + click/return for logical OR
Projects
Filter by project
Milestones
Filter by milestone
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

[Mosaic GPU] Prepare mamul example so it can be exposed to other projects.
#21255 opened May 16, 2024 by copybara-service bot Loading…
Add a sharding test for concatenate with an internal operand that requires
#21253 opened May 15, 2024 by copybara-service bot Loading…
[Pallas] Add temporary skip for pallas out-of-bounds interpret mode test on GPU.
#21252 opened May 15, 2024 by copybara-service bot Loading…
@justinjfu
Initial commit for jax.experimental.compute_on API.
#21249 opened May 15, 2024 by copybara-service bot Loading…
@yashk2810
[Mosaic GPU] Enable pointwise operation with scalar values.
#21248 opened May 15, 2024 by copybara-service bot Loading…
Add jax.tree_util.register_simple
#21245 opened May 15, 2024 by jakevdp • Draft
@jakevdp
Add core count to tpu nightly, fix v5 job
#21244 opened May 15, 2024 by MichaelHudgins Loading…
1
Add boundary='wrap' support for jax.scipy.signal.convolve2d and jax.scipy.signal.correlate2d
#21241 opened May 15, 2024 by rajasekharporeddy Loading…
1
@jakevdp
19
[DOC] Add documentation for jax.experimental.serialize_executable module
#21231 opened May 14, 2024 by nouiz Loading…
Test the loading of parameters from host memory
#21230 opened May 14, 2024 by copybara-service bot Loading…
Switch the order of sharding and memory kind custom call
#21229 opened May 14, 2024 by copybara-service bot Loading…
@yashk2810
Add new flash attn features to cuDNN SDPA API and remove fused attn
#21228 opened May 14, 2024 by Cjkkkk Loading…
@superbobry
1
[export] Enable model replication sharding using jax.pmap.
#21219 opened May 14, 2024 by copybara-service bot Loading…
Pulled build_mosaic_gpu default into jax.bzl
#21205 opened May 13, 2024 by copybara-service bot Loading…
@superbobry
Transpose for reduce window.
#21197 opened May 12, 2024 by copybara-service bot Loading…
1
[Mosaic GPU] Implement non-transposing matmul for bf16 x s8 matmul
#21194 opened May 12, 2024 by copybara-service bot Loading…
1
fix: scipy rel_entr edge case
#21193 opened May 12, 2024 by sh0416 Loading…
@jakevdp
16
Add Distributed data loading in a multi-host/multi-process environment doc
#21181 opened May 10, 2024 by 8bitmp3 Loading…
@jakevdp
2
Drop refs to australis in build rules and exporter.
#21178 opened May 10, 2024 by copybara-service bot Loading…
@pschuh
lobpcg: avoid changing X64 setting
pull ready Ready for copybara import and testing
#21165 opened May 10, 2024 by jakevdp Loading…
@jakevdp
Update Jupytext version to 1.16.1 in developer.md
#21164 opened May 10, 2024 by 8bitmp3 Loading…
Port JAX LAPACK kernels to XLA's FFI (API v4)
#21162 opened May 10, 2024 by copybara-service bot Loading…
1
XLA - Remove CreateToken, integ
#21156 opened May 10, 2024 by copybara-service bot Loading…
Bump version numbers post v0.4.28 release.
kokoro:force-run pull ready Ready for copybara import and testing
#21154 opened May 9, 2024 by hawkinsp Loading…
4
Add support to pipeline emitter for shapes that don't perfectly divide the block shapes
#21153 opened May 9, 2024 by copybara-service bot Loading…
@sharadmv
ProTip! Filter pull requests by the default branch with base:main.