Skip to content

Pull requests: google/jax

New pull request New
373 Open 13,856 Closed
373 Open 13,856 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

[export] Disable serialization in export_test if flatbuffers is not installed pull ready Ready for copybara import and testing
#22199 opened Jun 29, 2024 by gnecula Loading…
@gnecula
skip export_test.py if flatbuffers not importable
#22195 opened Jun 29, 2024 by mattjj Loading…
1
make pe.inline_jaxpr_into_trace work with dynamic shapes
#22194 opened Jun 29, 2024 by mattjj Loading…
Improve documentation of jax.numpy.clip documentation
#22193 opened Jun 28, 2024 by jakevdp Loading…
@jakevdp
better log10 doc
#22192 opened Jun 28, 2024 by pkgoogle Loading…
@dfm
1
better log2 doc
#22191 opened Jun 28, 2024 by pkgoogle Loading…
@dfm
2
[Pallas] Add vmap support to to_pallas_key
#22190 opened Jun 28, 2024 by copybara-service bot Loading…
@justinjfu
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
[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…
[Mosaic GPU] Pass in TMA descriptors through kernel parameters
#22175 opened Jun 28, 2024 by copybara-service bot Loading…
@apaszke
Don't use the trace context in the prune_closed_jaxpr_outputs cache.
#22174 opened Jun 28, 2024 by copybara-service bot Loading…
@hawkinsp
[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…
DOC: Improve docs of jnp.var and std
#22167 opened Jun 28, 2024 by rajasekharporeddy Loading…
@jakevdp
10
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
[pallas] Added more documentation for grid and BlockSpec. pull ready Ready for copybara import and testing
#22134 opened Jun 27, 2024 by gnecula Loading…
@gnecula
31
DOC: Update jnp.mean doc
#22132 opened Jun 27, 2024 by rajasekharporeddy Loading…
@jakevdp
4
Disable pallas_pipeline_tpu_test under asan/msan/tsan.
#22122 opened Jun 26, 2024 by copybara-service bot Loading…
@levskaya
Remove linear parameter from lax.cond_p pull ready Ready for copybara import and testing
#22119 opened Jun 26, 2024 by dfm Loading…
@dfm
2
Fix gpu_jax_head_jaxlib_pypi_latest job after migrating to plugin structure for jaxlib dependency
#22115 opened Jun 26, 2024 by copybara-service bot Loading…
[memories] Transfer to pinned_host fast path in async_serialize
#22114 opened Jun 26, 2024 by gspschmid Loading…
10
Demo export cl
#22112 opened Jun 26, 2024 by copybara-service bot Loading…
ProTip! Filter pull requests by the default branch with base:main.