Skip to content

Pull requests: google/jax

New pull request New
311 Open 13,210 Closed
311 Open 13,210 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

Simplify JAX lowering rules for cumulative sum
#21274 opened May 17, 2024 by copybara-service bot Loading…
Removed noop # type: ignore comments
pull ready Ready for copybara import and testing
#21273 opened May 17, 2024 by superbobry Loading…
1
[Mosaic] Also check bitwidth in apply-vector-layout's layoutIsValidForValue.
#21272 opened May 17, 2024 by copybara-service bot Loading…
[Mosaic GPU] A facility to check for performance regressions.
#21271 opened May 17, 2024 by copybara-service bot Loading…
1
[Mosaic GPU] Fix the accumulator sync fence by defining the registers instead of using constants.
#21270 opened May 17, 2024 by copybara-service bot Loading…
1
jax/pallas support ellipsis indexing
#21266 opened May 16, 2024 by copybara-service bot Loading…
1
jax.nn.softmax: only use initial=-inf when necessary
#21264 opened May 16, 2024 by copybara-service bot Loading…
@jakevdp
Deprecated ddof in jnp.var and jnp.std
pull ready Ready for copybara import and testing
#21262 opened May 16, 2024 by vfdev-5 • Draft
@jakevdp
10
CompatTest.test_cuda_eigh_cusolver_syev_dtype=f32_syevd: skip for CUDA 12.4+
#21258 opened May 16, 2024 by olupton Loading…
Make jax.grad and compute_on work correctly. If the forward pass has annotation to execute on CPU, then it's backward pass also executes on CPU.
#21257 opened May 16, 2024 by copybara-service bot Loading…
@yashk2810
[Mosaic GPU] Prepare mamul example so it can be exposed to other projects.
#21255 opened May 16, 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
Add jax.tree_util.register_simple
#21245 opened May 15, 2024 by jakevdp • Draft
@jakevdp
Add boundary='wrap' support for jax.scipy.signal.convolve2d and jax.scipy.signal.correlate2d
pull ready Ready for copybara import and testing
#21241 opened May 15, 2024 by rajasekharporeddy Loading…
1
@jakevdp
32
[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…
Add new flash attn features to cuDNN SDPA API and remove fused attn
#21228 opened May 14, 2024 by Cjkkkk Loading…
@superbobry
3
[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
18
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
ProTip! Type g p on any issue or pull request to go back to the pull request listing page.