Skip to content

Pull requests: jax-ml/jax

New pull request New
713 Open 27,139 Closed
713 Open 27,139 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
Assigned to nobody Loading
Sort
Sort by
Newest Oldest Most commented Least commented Recently updated Least recently updated Best match
Most reactions
👍 👎 😄 🎉 😕 ❤️ 🚀 👀

Pull requests list

[Mosaic TPU] Deprecated tpu.rotate. Use tpu.dynamic_rotate instead, which supports both static and dynamic rotate.
#38422 opened Jun 11, 2026 by copybara-service Bot Loading…
[Mosaic] Avoid MLIR canonicalizer pattern for comparisons that can result in unsupported data types.
#38420 opened Jun 11, 2026 by copybara-service Bot Loading…
take_current_trace should leave no trace kokoro:force-run pull ready Ready for copybara import and testing
#38415 opened Jun 11, 2026 by mattjj Collaborator Loading…
@mattjj
1
Add a jtu.is_libtpu_at_least version guard. kokoro:force-run pull ready Ready for copybara import and testing
#38407 opened Jun 11, 2026 by hawkinsp Collaborator Loading…
3
Make persistent compilation cache support host callbacks kokoro:force-run pull ready Ready for copybara import and testing
#38405 opened Jun 11, 2026 by skye-ant Loading…
@pschuh
3
Bump astral-sh/setup-uv from 7.3.1 to 8.2.0 dependencies Pull requests that update a dependency file github_actions Pull requests that update Github_actions code
#38402 opened Jun 11, 2026 by dependabot Bot Loading…
Bump aws-actions/configure-aws-credentials from 6.1.0 to 6.2.0 dependencies Pull requests that update a dependency file github_actions Pull requests that update Github_actions code
#38400 opened Jun 11, 2026 by dependabot Bot Loading…
Added debug arg to plgpu.kernel function
#38399 opened Jun 11, 2026 by vfdev-5 Member • Draft
3
Add a config flag to be able to blanket disabling of custom call safety checks in export.
#38393 opened Jun 11, 2026 by copybara-service Bot Loading…
Add bcomplex32 (complex<bfloat16>) ExtendedDType support
#38392 opened Jun 11, 2026 by DeanTMaxim Loading…
4 of 5 tasks
@jakevdp
6
Propagate NaNs in uniform stats functions
#38390 opened Jun 11, 2026 by KirtiRamchandani Loading…
1
Fix istft noverlap and inferred nfft handling
#38389 opened Jun 11, 2026 by KirtiRamchandani Loading…
1
2
Removed another absl.logging usage in pallas codebase
#38388 opened Jun 11, 2026 by vfdev-5 Member • Draft
[doc] Add some hints for debugging and fixing slow tracing and compilation times
#38387 opened Jun 11, 2026 by copybara-service Bot Loading…
@gnecula
[Mosaic] Remove temporary backwards compatibility builder for SC vector core with null subcore index.
#38381 opened Jun 10, 2026 by copybara-service Bot Loading…
@emilyfertig
Reverts 5fffc8c8002b33cb8d01d5c9973335a1b0aadd64
#38377 opened Jun 10, 2026 by copybara-service Bot Loading…
[JAX] Add py_pjrt_host_callback extension to jaxlib
#38374 opened Jun 10, 2026 by copybara-service Bot Loading…
Add Device.clear_memory_stats()
#38373 opened Jun 10, 2026 by fl0pedro Loading…
3
Refactor JAX config and name stack decorators to be thread-safe and re-entrant.
#38365 opened Jun 10, 2026 by copybara-service Bot Loading…
Use vars instead of argument-list indices in effect types
#38362 opened Jun 10, 2026 by copybara-service Bot Loading…
@dougalm
docs: improve grammar and wording in Ref guide
#38361 opened Jun 10, 2026 by dhruvatr Loading…
1
Add faster approximate gamma sampler
#38358 opened Jun 10, 2026 by Gattocrucco Contributor Loading…
1
[NFC] [pallas] A very rough take on revamping the quickstart documentation
#38350 opened Jun 10, 2026 by superbobry Collaborator • Draft
18
Wrap negative indices by default in take_along_axis and remove legacy behavior that used promise_in_bounds.
#38341 opened Jun 10, 2026 by copybara-service Bot Loading…
internal change
#38334 opened Jun 10, 2026 by copybara-service Bot Loading…
ProTip! Exclude everything labeled bug with -label:bug.