Skip to content

Pull requests: google/jax

New pull request New
362 Open 13,903 Closed
362 Open 13,903 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

make remat reduce precision of saved values to avoid xla excess precision
#22244 opened Jul 2, 2024 by mattjj • Draft
Extend CustomCallOp backend_config to take a DictionaryAttr
#22242 opened Jul 2, 2024 by copybara-service bot Loading…
jax.numpy: better docstring for array_equiv function
#22241 opened Jul 2, 2024 by selamw1 Loading…
Add the support of attention logits soft capping to the PagedAttention Pallas kernel.
#22239 opened Jul 2, 2024 by copybara-service bot Loading…
Prefer __qualname__ as a pjit_p name pull ready Ready for copybara import and testing
#22237 opened Jul 2, 2024 by hawkinsp Loading…
1
[Pallas] Add better reduction support.
#22236 opened Jul 2, 2024 by copybara-service bot Loading…
@justinjfu
Enable opt-in autodetection of distributed configuration for mpi4py, attempt 2
#22235 opened Jul 2, 2024 by coreyjadams • Draft
2
[Mosaic_GPU] Allow measured functions to be rerun without recompiling
#22234 opened Jul 2, 2024 by copybara-service bot Loading…
[xla:ffi] Use lazy decoding for Buffer<dtype,rank>
#22224 opened Jul 1, 2024 by copybara-service bot Loading…
Introduce a canonicalize pass, rewrite all contractions as matmuls (vector::ContractionOp as tpu::MatMulOp), remove special handling for contraction op in other passes.
#22220 opened Jul 1, 2024 by copybara-service bot Loading…
Don't wrap singleton ir.Values with tuples during HLO lowering. pull ready Ready for copybara import and testing
#22211 opened Jul 1, 2024 by hawkinsp Loading…
1
Improve documentation of jax.numpy.clip documentation pull ready Ready for copybara import and testing
#22193 opened Jun 28, 2024 by jakevdp Loading…
@jakevdp
1
better log10 doc pull ready Ready for copybara import and testing
#22192 opened Jun 28, 2024 by pkgoogle Loading…
@dfm
4
better log2 doc pull ready Ready for copybara import and testing
#22191 opened Jun 28, 2024 by pkgoogle Loading…
@dfm
5
Better docs for jnp.sort & jnp.argsort documentation pull ready Ready for copybara import and testing
#22189 opened Jun 28, 2024 by jakevdp Loading…
@jakevdp
2
[Mosaic] Column shift relayouts for non-native tilings and packed types, except for (1, n) and packed
#22188 opened Jun 28, 2024 by copybara-service bot Loading…
[pallas] Improve some error messages and add API tests. pull ready Ready for copybara import and testing
#22173 opened Jun 28, 2024 by gnecula Loading…
@gnecula
[Mosaic GPU] Support narrower swizzles in copy and TMA tests
#22172 opened Jun 28, 2024 by copybara-service bot Loading…
@apaszke
[memories] Overlap shard transfers in async_serialize
#22169 opened Jun 28, 2024 by gspschmid Loading…
jnp.arange: fix incorrect type annotation pull ready Ready for copybara import and testing
#22160 opened Jun 28, 2024 by jakevdp Loading…
@jakevdp
Update CUDA custom call example code to use ffi_call
#22141 opened Jun 27, 2024 by dfm • Draft
@dfm
Improve docs and add test case for RegularGridInterpolator
#22138 opened Jun 27, 2024 by timmens Loading…
@jakevdp
[memories] Transfer to pinned_host fast path in async_serialize
#22114 opened Jun 26, 2024 by gspschmid Loading…
11
Demo export cl
#22112 opened Jun 26, 2024 by copybara-service bot Loading…
[WIP] Add a wrapper class for etils.epath
#22108 opened Jun 26, 2024 by ayaka14732 • Draft
ProTip! Adding no:label will show everything without a label.