Skip to content

Pull requests: google/jax

New pull request New
344 Open 13,588 Closed
344 Open 13,588 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

Compilation cache doc
#21819 opened Jun 12, 2024 by keshavb96 • Draft
Add frontend attributes to Jax. This allows Jax users to annotate Jax code with frontend_attributes which can be traced down to the HLO level, to be used for numerical debugging purposes. Frontend attributes annotations for instructions added by grad are not supported yet.
#21818 opened Jun 12, 2024 by copybara-service bot Loading…
[PJRT:PLUGIN] Support both string and bytes as the input type of function name for register_custom_call_target in jax-cuda-plugin.
#21814 opened Jun 11, 2024 by copybara-service bot Loading…
@jyingl3
Enable jax's cloud-tpu configs when libtpu is present via through "pip install" or set by custom through the $TPU_LIBRARY_PATH env var
#21812 opened Jun 11, 2024 by copybara-service bot Loading…
Add freshness metablock to JAX OSS docs.
#21810 opened Jun 11, 2024 by copybara-service bot Loading…
[ROCm]: Pallas updates for ROCm
#21809 opened Jun 11, 2024 by rahulbatra85 Loading…
1
[Mosaic GPU] Change FlashAttention implementation to support Grouped Query Attention.
#21808 opened Jun 11, 2024 by copybara-service bot Loading…
@bchetioui
prune cond passthrough outputs pull ready Ready for copybara import and testing
#21806 opened Jun 11, 2024 by cgarciae Loading…
DOC: Better docs for jnp.convolve & jnp.correlate
#21804 opened Jun 11, 2024 by rajasekharporeddy Loading…
@jakevdp
9
[pallas] Add support for cross-platform lowering pull ready Ready for copybara import and testing
#21799 opened Jun 11, 2024 by gnecula Loading…
@gnecula
16
doc: don't use typehints for signatures. documentation
#21797 opened Jun 11, 2024 by jakevdp Loading…
@jakevdp
5
[Mosaic] Simplify vector.shape_cast rules and cover more cases
#21796 opened Jun 11, 2024 by copybara-service bot Loading…
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…
[export] Add documentation for jax.export pull ready Ready for copybara import and testing
#21769 opened Jun 10, 2024 by gnecula Loading…
@gnecula
10
Improve vmap out_axes error
#21766 opened Jun 10, 2024 by cgarciae Loading…
3
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…
Activate Cholesky Factorization to XLA's FFI
#21727 opened Jun 7, 2024 by copybara-service bot Loading…
1
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
ProTip! Type g i on any issue or pull request to go back to the issue listing page.