Skip to content

Pull requests: google/jax

New pull request New
348 Open 13,561 Closed
348 Open 13,561 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

Remove unused variable pull ready Ready for copybara import and testing
#21791 opened Jun 10, 2024 by jakevdp Loading…
@jakevdp
DOC: jax.lax.top_k: fix docstring rendering & add example documentation pull ready Ready for copybara import and testing
#21788 opened Jun 10, 2024 by jakevdp Loading…
@jakevdp
Simplify definition of jax.scipy.special.kl_div
#21782 opened Jun 10, 2024 by jakevdp Loading…
@jakevdp
2
Pallas bitwise_left_shift unit test fix pull ready Ready for copybara import and testing
#21780 opened Jun 10, 2024 by rahulbatra85 Loading…
3
Rewrite vector.multi_dim_reduction with bf16 source/accumulator/output into
#21775 opened Jun 10, 2024 by copybara-service bot Loading…
[Pallas] Add pallas TPU random key impls and lowering rules for basic prng ops (seed/foldin/bits/unwrap/wrap).
#21773 opened Jun 10, 2024 by copybara-service bot Loading…
@justinjfu
jax.scipy.special.beta: deprecate x,y in favor of a,b
#21772 opened Jun 10, 2024 by jakevdp Loading…
@jakevdp
[export] Add documentation for jax.export
#21769 opened Jun 10, 2024 by gnecula Loading…
5
Improve vmap out_axes error
#21766 opened Jun 10, 2024 by cgarciae Loading…
1
DOC: Add example code snippets to jax.scipy.signal.convolve2d, correlate and correlate2d docs
#21762 opened Jun 10, 2024 by rajasekharporeddy Loading…
Batch pxla.shard_args calls triggered by jax.device_put
#21754 opened Jun 9, 2024 by copybara-service bot Loading…
Make full Read the Docs builds for PR previews explicitly opt-in
#21750 opened Jun 9, 2024 by dfm • Draft
8
Allow specifying initial forget gate bias for LSTMs.
#21749 opened Jun 9, 2024 by gauchm Loading…
Add ffi_call function with a similar signature to pure_callback.
#21743 opened Jun 7, 2024 by copybara-service bot Loading…
@dfm
Enable runtime uptime telemetry for JAX on Cloud TPU.
#21737 opened Jun 7, 2024 by copybara-service bot Loading…
Make pxla.shard_arg batch calls to xc.copy_array_to_devices_with_sharding
#21736 opened Jun 7, 2024 by copybara-service bot Loading…
Updated the type annotations of *_spec= parameters of pl.pallas_call pull ready Ready for copybara import and testing
#21729 opened Jun 7, 2024 by superbobry Loading…
Removed kernel_regeneration_util from Mosaic
#21728 opened Jun 7, 2024 by copybara-service bot Loading…
@superbobry
Activate Cholesky Factorization to XLA's FFI
#21727 opened Jun 7, 2024 by copybara-service bot Loading…
1
DOC: Add example code snippets to jax.scipy.linalg.expm and jax.scipy.linalg.polar docs
#21725 opened Jun 7, 2024 by rajasekharporeddy Loading…
10
Fix a bug where treedef.flatten_up_to(...) was overly permissive for None treedefs.
#21709 opened Jun 6, 2024 by copybara-service bot Loading…
@hawkinsp
Passes kv_store to ifrt::PjRtClient.
#21707 opened Jun 6, 2024 by copybara-service bot Loading…
@jyingl3
[Mosaic GPU] Satisfy MLIR that the dialect interfaces are registered to fix UNDEBUG
#21698 opened Jun 6, 2024 by copybara-service bot Loading…
Migrates multihost_utils to use make_array_from_process_local_data.
#21693 opened Jun 6, 2024 by copybara-service bot Loading…
Refactor & Document
#21691 opened Jun 6, 2024 by Arkay92 Loading…
2
ProTip! Follow long discussions with comments:>50.