Skip to content
New pull request New
174 Open 5,915 Closed
174 Open 5,915 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
Most reactions
👍 👎 😄 🎉 😕 ❤️ 🚀 👀
Roll forward again after the fix in the auto sharding pass.
#10127 opened Apr 1, 2022 by copybara-service bot
Remove the tuple when there is only one axis of the mesh mapped to a dimension of the input. Improves the readability.
#10125 opened Apr 1, 2022 by copybara-service bot
prototyping dynamic shapes
pull ready
#10120 opened Apr 1, 2022 by mattjj
2
* Disallow any other type other than GDA and ShapedArray for auto sharding.
#10111 opened Mar 31, 2022 by copybara-service bot
@yashk2810
[JAX] Adds the approx_top_k_p bridge.
#10108 opened Mar 31, 2022 by copybara-service bot
1
Remove flags set for all v4 TPUs. Topology flags will now be set in libTPU.
#10098 opened Mar 31, 2022 by copybara-service bot
1
Remove the jax_enable_mlir flag. MLIR is now the only supported code path.
#10093 opened Mar 30, 2022 by copybara-service bot
@hawkinsp
Re-land changes from https://github.com/google/jax/pull/10069
#10088 opened Mar 30, 2022 by copybara-service bot
@jakevdp
Expose CopyToRemoteDevice and MakeCrossHostReceiveBuffer in Python bindings.
#10086 opened Mar 30, 2022 by copybara-service bot
Allow the same bitwise quantization types in MHLO ADD op
#10040 opened Mar 28, 2022 by copybara-service bot
Add scipy.linalg.rsf2csf
pull ready
#10039 opened Mar 27, 2022 by ajcr • Approved
@hawkinsp
2
use weakref_lru_cache from #9966 in more places
#10036 opened Mar 25, 2022 by mattjj • Draft
@mattjj
2
Migrate away from using JaxTestCase in tests
#10024 opened Mar 24, 2022 by copybara-service bot
@jakevdp
WIP: Try enabling JAX experimental name stack
#10015 opened Mar 23, 2022 by copybara-service bot
@sharadmv
generic protocol for compilation wrappers
#9999 opened Mar 23, 2022 by YouJiacheng
1
@froystig
2
Enable use of GlobalDeviceArray (GDA) in T5X Checkpointer. Add a separate unit test, gda_checkpoints_test, to cover this use case.
#9990 opened Mar 22, 2022 by copybara-service bot
@yashk2810
[PJRT] Add CompiledMemoryStats and allow us to retrieve static memory usage for compiled programs in JAX.
#9983 opened Mar 21, 2022 by copybara-service bot
@zhangqiaorjc
fix jnp.convolve int precision
pull ready
#9953 opened Mar 18, 2022 by mattjj • Approved
@mattjj
1
elide trivial slices, cf. #9215
pull ready
#9926 opened Mar 16, 2022 by mattjj
@mattjj
2
Safe where in jax.nn.celu and nn.leaky_relu
#9876 opened Mar 14, 2022 by VolodyaCO
1
@jakevdp
4
Fix _lax_max_taylor_rule for broadcastable inputs
#9868 opened Mar 14, 2022 by YouJiacheng
1
add use_auto_spmd_partitioning to the rest of places wherever use_spmd_partitioning exists
#9822 opened Mar 9, 2022 by copybara-service bot
Switch gpu, cpu and jax2tf to use the new OptimizationBarrier op.
#9809 opened Mar 9, 2022 by pschuh
@pschuh
2
Replace an assertion with no error message with an error message.
#9797 opened Mar 8, 2022 by jblespiau
@mattjj
1
Adds use_auto_spmd_partitioning and propagate its value down the stack
#9790 opened Mar 7, 2022 by copybara-service bot
ProTip! Find all pull requests that aren't related to any open issues with -linked:issue.