Skip to content

Pull requests: google/jax

New pull request New
314 Open 13,220 Closed
314 Open 13,220 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

Automated Code Change
#21294 opened May 18, 2024 by copybara-service bot Loading…
[ROCm] Cholesky update kernel fixes for ROCm
#21292 opened May 18, 2024 by Ruturaj4 Loading…
2
Add unit test to check if the backend serialization/deserialization result equal to the original executable.
#21291 opened May 17, 2024 by copybara-service bot Loading…
Finalize deprecation of the config module.
#21288 opened May 17, 2024 by copybara-service bot Loading…
@jakevdp
fix: Update hermetic python dependencies to numpy=2.0.0rc2 and scipy=1.13.0 for all python version
#21285 opened May 17, 2024 by copybara-service bot Loading…
[export] Relax the check that exported modules are used with same number of devices as when exported
pull ready Ready for copybara import and testing
#21280 opened May 17, 2024 by gnecula Loading…
@gnecula
Order axis names for shard_map residuals
#21278 opened May 17, 2024 by jaro-sevcik Loading…
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…
@superbobry
3
[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.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
12
CompatTest.test_cuda_eigh_cusolver_syev_dtype=f32_syevd: skip for CUDA 12.4+
#21258 opened May 16, 2024 by olupton Loading…
[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
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
pull ready Ready for copybara import and testing
#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
pull ready Ready for copybara import and testing
#21228 opened May 14, 2024 by Cjkkkk Loading…
@superbobry
12
[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
ProTip! Add no:assignee to see everything that’s not assigned.