Skip to content

Pull requests: jax-ml/jax

New pull request New
450 Open 18,183 Closed
450 Open 18,183 Closed
Author
Filter by author
Loading
Label
Filter by label
Loading
Use alt + click/return to exclude labels
or ⇧ + click/return for logical OR
Projects
Filter by project
Loading
Milestones
Filter by milestone
Loading
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

Merge tuple_replace and tuple_update in jax._src.util.
#27853 opened Apr 9, 2025 by carlosgmartin Loading…
[pallas:mosaic_gpu] ModuleContext.reserve_barrier is now a context manager
#27851 opened Apr 8, 2025 by copybara-service bot Loading…
@superbobry
Automatically document jax.config options
#27849 opened Apr 8, 2025 by ZacCranko • Draft
Remove redundant filtering in the ragged flash attention kernel
#27848 opened Apr 8, 2025 by copybara-service bot Loading…
Add unit tests for the grouped query attention reference implementation
#27845 opened Apr 8, 2025 by copybara-service bot Loading…
Remove deprecated support for jax.dlpack.from_dlpack accepting a capsule.
#27844 opened Apr 8, 2025 by copybara-service bot Loading…
@hawkinsp
Add JVP rule for linear_call kokoro:force-run pull ready Ready for copybara import and testing
#27843 opened Apr 8, 2025 by dfm Loading…
@dfm
Accelerate deprecation of vectorized argument in callbacks.
#27841 opened Apr 8, 2025 by copybara-service bot Loading…
@dfm
InspectSharding doesn't work for remote backends.
#27839 opened Apr 8, 2025 by copybara-service bot Loading…
Update debug_info tests for use_direct_linearize.
#27837 opened Apr 8, 2025 by copybara-service bot Loading…
@danielsuo
Fix linear_call to allow recursive definitions pull ready Ready for copybara import and testing
#27831 opened Apr 8, 2025 by dfm Loading…
@dfm
2
[JAX] Remove deprecated exports in jax.lib.xla_client.
#27826 opened Apr 8, 2025 by copybara-service bot Loading…
@hawkinsp
document SPMD pipeline parallelism
#27816 opened Apr 8, 2025 by copybara-service bot Loading…
[Pallas] Fix ()-shaped vectors being materialized in Pallas lowering.
#27811 opened Apr 7, 2025 by copybara-service bot Loading…
@justinjfu
#sdy Add more debug info when there is a mesh mismatch in JAX export
#27810 opened Apr 7, 2025 by copybara-service bot Loading…
[jaxlib] Pack/unpack subbyte types to/from numpy arrays to support int2, uint2, int4, uint4, float4_e2m1fn subbyte types in CPU/GPU callbacks.
#27806 opened Apr 7, 2025 by copybara-service bot Loading…
@danielsuo
[Pallas/Fuser] Add basic closed over consts support to pull_block_spec
#27801 opened Apr 7, 2025 by copybara-service bot Loading…
@sharadmv
roll back test-breaking change
#27788 opened Apr 7, 2025 by copybara-service bot Loading…
Add option for debug print to be called on partitioned arguments rather than all-gather and print the entire logical arguments.
#27781 opened Apr 7, 2025 by copybara-service bot Loading…
@danielsuo
[state] Added some docs to DischargeRule and PartialDischargeRule
#27772 opened Apr 7, 2025 by copybara-service bot Loading…
@superbobry
[Mosaic GPU] Add a LayoutCast op to the Mosaic GPU mlir dialect.
#27771 opened Apr 7, 2025 by copybara-service bot Loading…
[vmappable] fix trace context bugs kokoro:force-run pull ready Ready for copybara import and testing
#27759 opened Apr 6, 2025 by mattjj Loading…
Move l_next (float32) division out of the flash attention kernel
#27746 opened Apr 4, 2025 by copybara-service bot Loading…
Create a PjRt protos dir.
#27744 opened Apr 4, 2025 by copybara-service bot Loading…
Fix jax.random.randint for 8-bit and 16-bit dtypes kokoro:force-run pull ready Ready for copybara import and testing
#27742 opened Apr 4, 2025 by jakevdp Loading…
@jakevdp
4
ProTip! Find all pull requests that aren't related to any open issues with -linked:issue.