Skip to content

Pull requests: jax-ml/jax

New pull request New
467 Open 18,335 Closed
467 Open 18,335 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

[Draft] cudnn sdpa flex attention
#28048 opened Apr 15, 2025 by Cjkkkk • Draft
Clean up softmax initial deprecation pull ready Ready for copybara import and testing
#28047 opened Apr 15, 2025 by jakevdp Loading…
@jakevdp
convert NumPy RNG key data to uncommitted CPU-backed jax.Array data
#28045 opened Apr 15, 2025 by copybara-service bot Loading…
@froystig
Make Clang use manylinux C++ standard library
#28041 opened Apr 15, 2025 by charleshofer Loading…
Replicate the failure
#28038 opened Apr 15, 2025 by MichaelHudgins • Draft
Debugging CI failure.
#28035 opened Apr 15, 2025 by hawkinsp • Draft
Add jax_logging_level to common config options in docs documentation kokoro:force-run pull ready Ready for copybara import and testing
#28034 opened Apr 15, 2025 by skye Loading…
4
breaking internal tests
#28033 opened Apr 15, 2025 by copybara-service bot Loading…
Export mlir.custom_call from under extend.mlir.
#28031 opened Apr 15, 2025 by copybara-service bot Loading…
@dfm
[Pallas][Mosaic][TPU] Add disable_bounds_checks compiler params
#28028 opened Apr 15, 2025 by copybara-service bot Loading…
[Pallas Fuser] Add basic sublane/lane reshape fusion
#28026 opened Apr 15, 2025 by copybara-service bot Loading…
@sharadmv
Use Docker images with pre-installed Python.
#28024 opened Apr 15, 2025 by copybara-service bot Loading…
@MichaelHudgins
Initial f32, bf16 TPU SublanesShuffleOp
#28019 opened Apr 14, 2025 by copybara-service bot Loading…
disable cudnn attention packed layout test on old arch prior to hopper
#28018 opened Apr 14, 2025 by Cjkkkk Loading…
jax.jit: support decorator factory pattern kokoro:force-run pull ready Ready for copybara import and testing
#28011 opened Apr 14, 2025 by jakevdp Loading…
@jakevdp
9
Remove roofline subtests so they work with pytest.
#28010 opened Apr 14, 2025 by copybara-service bot Loading…
Added new multi-controller JAX guide. kokoro:force-run pull ready Ready for copybara import and testing
#28009 opened Apr 14, 2025 by mwhittaker Loading…
@mwhittaker
11
[WIP] Amend documentation of JAX config options, include context managers, mechanism to hide flags documentation
#28006 opened Apr 14, 2025 by ZacCranko • Draft
@ZacCranko
5
Add basic documentation for Pallas:MGPU layouts and transforms documentation pull ready Ready for copybara import and testing
#28004 opened Apr 14, 2025 by apaszke Loading…
@apaszke
Make ShardMapTracer track vma instead of rep.
#28002 opened Apr 14, 2025 by copybara-service bot Loading…
@yashk2810
Remove jax_varying_axes_in_types config and rewrite from shard_map_p
#28001 opened Apr 14, 2025 by copybara-service bot Loading…
@yashk2810
Update ManualComputation round-tripping to not have round-tripping attrs on the CallOps.
#28000 opened Apr 14, 2025 by copybara-service bot Loading…
Fix new_reference / _Py_IsOwnedByCurrentThread race suppression for 3.13
#27994 opened Apr 14, 2025 by vfdev-5 Loading…
Add extension point for letting TestLoader specify a custom sharding scheme.
#27987 opened Apr 13, 2025 by copybara-service bot Loading…
Simplify nn.log_softmax and add edit_excluded argument.
#27986 opened Apr 13, 2025 by carlosgmartin Loading…
1
ProTip! Mix and match filters to narrow down what you’re looking for.