Skip to content

Pull requests: google/jax

New pull request New
343 Open 13,605 Closed
343 Open 13,605 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 test for QDWH with dynamic shapes.
#21841 opened Jun 13, 2024 by copybara-service bot Loading…
avoid globally silencing the jit backend/device warning
#21834 opened Jun 12, 2024 by jakevdp Loading…
@jakevdp
Deprecate internal pretty-printing APIs, jax.core.pp_* pull ready Ready for copybara import and testing
#21829 opened Jun 12, 2024 by jakevdp Loading…
@jakevdp
[export] Rename some API entry points pull ready Ready for copybara import and testing
#21828 opened Jun 12, 2024 by gnecula Loading…
@gnecula
9
Adding GCS support for LRUCache
#21822 opened Jun 12, 2024 by ayaka14732 • Draft
3 tasks
@ayaka14732
Compilation cache doc
#21819 opened Jun 12, 2024 by keshavb96 Loading…
@ayaka14732
3
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…
Add freshness metablock to JAX OSS docs.
#21810 opened Jun 11, 2024 by copybara-service bot Loading…
[ROCm]: Pallas updates for ROCm pull ready Ready for copybara import and testing
#21809 opened Jun 11, 2024 by rahulbatra85 Loading…
5
prune cond passthrough outputs pull ready Ready for copybara import and testing
#21806 opened Jun 11, 2024 by cgarciae Loading…
1
DOC: Better docs for jnp.convolve & jnp.correlate
#21804 opened Jun 11, 2024 by rajasekharporeddy Loading…
@jakevdp
19
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…
4
Rewrite vector.multi_dim_reduction with bf16 source/accumulator/output into
#21775 opened Jun 10, 2024 by copybara-service bot Loading…
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
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
Keeps kv_store in ifrt::PjRtClient if it is set in the CreateOption.
#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…
ProTip! Adding no:label will show everything without a label.