Skip to content

Pull requests: jax-ml/jax

New pull request New
439 Open 17,386 Closed
439 Open 17,386 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

Add use_high_dynamic_range_gumbel flag which allows sampling gumbel such
#26870 opened Mar 1, 2025 by copybara-service bot Loading…
@pschuh
[Pallas] Adds a simple dynamic race detector for TPU interpret mode.
#26866 opened Feb 28, 2025 by copybara-service bot Loading…
@jburnim
Improved errors when indexing with floats better_errors Improve the error reporting kokoro:force-run pull ready Ready for copybara import and testing
#26865 opened Feb 28, 2025 by jakevdp Loading…
@jakevdp
1
Implement Jax CPU/GPU callbacks with XLA's FFI.
#26864 opened Feb 28, 2025 by copybara-service bot Loading…
@danielsuo
doc: fix description of logsumexp axis documentation pull ready Ready for copybara import and testing
#26862 opened Feb 28, 2025 by jakevdp Loading…
@jakevdp
1
Delete ParsedPartitionSpec and preprocess function and couple more cleanups
#26861 opened Feb 28, 2025 by copybara-service bot Loading…
@yashk2810
Fix FFI callbacks for OSS.
#26857 opened Feb 28, 2025 by danielsuo Loading…
@danielsuo
[Pallas/Fuser] Add support for custom_call_jvp/pjit to push_block_spec
#26856 opened Feb 28, 2025 by copybara-service bot Loading…
@sharadmv
Add fuser to jax.experimental.pallas
#26855 opened Feb 28, 2025 by copybara-service bot Loading…
@sharadmv
[Pallas] Add experimental (private for now) API for manual fusion into Pallas kernels
#26854 opened Feb 28, 2025 by copybara-service bot Loading…
@sharadmv
Add scalar event logging function
#26853 opened Feb 28, 2025 by jeffcarp Loading…
Proof-of-concept: A JAX callback that operates directly on device buffers.
#26850 opened Feb 28, 2025 by copybara-service bot Loading…
@dfm
[Mosaic GPU][NFC] Move the calculation of group strides into _validate_mma
#26849 opened Feb 28, 2025 by copybara-service bot Loading…
@apaszke
[Mosaic GPU][NFC] Clean up the computation of group strides
#26848 opened Feb 28, 2025 by copybara-service bot Loading…
@apaszke
[mgpu] Non-vector untiled stores for tiling layouts.
#26847 opened Feb 28, 2025 by copybara-service bot Loading…
[Pallas/Mosaic GPU] Add lowerings/layout inference for all the necessary conversion ops when using Warpgroup semantics.
#26844 opened Feb 28, 2025 by copybara-service bot Loading…
@bchetioui
[mgpu] Forach in tiled layout.
#26842 opened Feb 28, 2025 by copybara-service bot Loading…
DOC: Remove duplicate JAX version 0.4.37 heading in Changelog
#26840 opened Feb 28, 2025 by rajasekharporeddy Loading…
Fixed printing order of results in jax.debug.print documentation
#26839 opened Feb 28, 2025 by Sai-Suraj-27 Loading…
1
Plumb layout through the creation of IFRT Arrays (roll-forward with fix).
#26838 opened Feb 28, 2025 by copybara-service bot Loading…
@emilyfertig
Fix thread safety of JAX error checking
#26837 opened Feb 28, 2025 by copybara-service bot Loading…
@ayaka14732
Added butterworth filter code
#26824 opened Feb 27, 2025 by ParamThakkar123 Loading…
1
Update ad.backward_pass to support non-linear functions of constants
#26811 opened Feb 27, 2025 by dfm Loading…
@dfm
6
Fix assertion fails on DebugInfo with JAX_USE_DIRECT_LINEARIZE=1
#26810 opened Feb 27, 2025 by gnecula Loading…
1
Remove tensorstore dependency from test dependencies for python3.13-nogil jobs.
#26777 opened Feb 26, 2025 by copybara-service bot Loading…
ProTip! Adding no:label will show everything without a label.