Skip to content

Pull requests: google/jax

New pull request New
310 Open 13,230 Closed
310 Open 13,230 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

Added a primitive for jax.debug.print
#21304 opened May 20, 2024 by copybara-service bot Loading…
@superbobry
[pallas] align interpreter load/store with masked behaviour
#21298 opened May 18, 2024 by oliverdutton Loading…
1
Automated Code Change
#21294 opened May 18, 2024 by copybara-service bot Loading…
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…
[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
14
CompatTest.test_cuda_eigh_cusolver_syev_dtype=f32_syevd: skip for CUDA 12.4+
#21258 opened May 16, 2024 by olupton 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
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
18
[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
ProTip! Find all pull requests that aren't related to any open issues with -linked:issue.