Skip to content

Pull requests: google/jax

New pull request New
348 Open 10,336 Closed
348 Open 10,336 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

Add Rotation functions
pull ready Ready for copybara import and testing
#17192 opened Aug 19, 2023 by chrisflesher Loading…
1
Add docs for is_fully_addressable to jax.Array and remove GDA from process_allgather docs and clarify it a bit more.
pull ready Ready for copybara import and testing
#17189 opened Aug 19, 2023 by copybara-service bot Loading…
@yashk2810
1
jnp.piecewise: don't cast outputs to input dtype
pull ready Ready for copybara import and testing
#17186 opened Aug 18, 2023 by jakevdp • Draft
@jakevdp
1
Test idempotent serialization of proto.
#17179 opened Aug 18, 2023 by copybara-service bot Loading…
1
[PJRT C API] Support passing allocator_config as an option in PJRT GPU plugin.
#17177 opened Aug 18, 2023 by copybara-service bot Loading…
1
Reverts c6f40e202c7f5724b9be61afa33541a8f4abfdd0
#17161 opened Aug 17, 2023 by copybara-service bot Loading…
@jakevdp
[Pallas] Add Mosaic lowering for pow2
#17148 opened Aug 16, 2023 by copybara-service bot Loading…
@sharadmv
create jax.extend.random as a copy of jax.prng
#17139 opened Aug 16, 2023 by copybara-service bot Loading…
@froystig
[host_callback] Change host_callback lowering to use custom calls on GPU.
#17137 opened Aug 16, 2023 by copybara-service bot Loading…
@gnecula
Ensure that CompileOptions serializes deterministically.
#17133 opened Aug 16, 2023 by copybara-service bot Loading…
1
[MOSAIC] apply_vector_layout rewrite (6): arith.constant
#17127 opened Aug 15, 2023 by copybara-service bot Loading…
[MOSAIC] apply_vector_layout C++ rewrite (5) vector load and store
#17125 opened Aug 15, 2023 by copybara-service bot Loading…
[JAX] Switch host_callback to use MLIR lowering instead of the older direct HLO translation rules.
#17124 opened Aug 15, 2023 by copybara-service bot Loading…
@hawkinsp
Remove all testing of jax.experimental.sparse on cusparse
#17122 opened Aug 15, 2023 by jakevdp Loading…
1
@jakevdp
Load Python rules from rules_python
#17119 opened Aug 15, 2023 by copybara-service bot Loading…
jnp.ufunc: add fast paths for add/prod reductions
#17113 opened Aug 14, 2023 by jakevdp Loading…
@jakevdp
3
[shard-map] add rewrite for efficient transposition
#17111 opened Aug 14, 2023 by mattjj • Draft
@mattjj
Create a context manager to register/unregister event duration listeners.
#17102 opened Aug 14, 2023 by copybara-service bot Loading…
1
[jax_triton] Always use "triton_kernel_call" as custom call target name.
#17099 opened Aug 14, 2023 by copybara-service bot Loading…
@chr1sj0nes
Add register_static decorator to tree_util to facilitate creating leafless classes.
#17096 opened Aug 14, 2023 by copybara-service bot Loading…
@hawkinsp
1
Added vectorized arguments to debug.{print,breakpoint,callback}.
#17088 opened Aug 12, 2023 by patrick-kidger Loading…
New cache key generation: CompileOptions.
#17083 opened Aug 11, 2023 by copybara-service bot Loading…
1
Make searchsorted a primitive
pull ready Ready for copybara import and testing
#17079 opened Aug 11, 2023 by jakevdp • Draft
@jakevdp
[Mosaic] Use strided load to load one entire row more efficiently.
#17073 opened Aug 11, 2023 by copybara-service bot Loading…
lax.expm1: more accurate gradient for large negative inputs
pull ready Ready for copybara import and testing
#17067 opened Aug 10, 2023 by jakevdp Loading…
@jakevdp
1
ProTip! Follow long discussions with comments:>50.