Skip to content

Pull requests: jax-ml/jax

New pull request New
514 Open 19,121 Closed
514 Open 19,121 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

[pallas] Added pl.loop -- a decorator for writing stateless loops
#29041 opened May 27, 2025 by copybara-service bot Loading…
@superbobry
[pallas:mosaic_gpu] Added the missing resource estimation rule for pl.run_state
#29040 opened May 27, 2025 by copybara-service bot Loading…
@superbobry
[tokamax] Persistent ragged dot pallas mgpu kernel (via nd_loop)
#29039 opened May 27, 2025 by copybara-service bot Loading…
[tokamax] Serialization supports RaggedDotDimensionNumbers.
#29038 opened May 27, 2025 by copybara-service bot Loading…
[pallas] Use the mgpu barrier rather than the dialect one in run_scoped lowering.
#29036 opened May 27, 2025 by copybara-service bot Loading…
[Mosaic GPU] Bug reproducer
#29034 opened May 27, 2025 by copybara-service bot Loading…
@apaszke
#sdy Fallback to GSPMD in JAX export if the loaded module was lowered for GSPMD.
#29033 opened May 27, 2025 by copybara-service bot Loading…
Fix handling of input None in custom_transpose.
#29030 opened May 27, 2025 by dfm Loading…
1
@dfm
[Mosaic GPU] Make the Blackwell matmul kernel persistent
#29029 opened May 27, 2025 by copybara-service bot Loading…
@apaszke
[pallas:triton] Removed the Triton prefix from TritonCompilerParams
#29023 opened May 27, 2025 by copybara-service bot Loading…
@superbobry
[mosaic_gpu] Use DIScopeForLLVMFuncOpPass from MLIR instead of its Triton fork
#29022 opened May 27, 2025 by copybara-service bot Loading…
@superbobry
Automated Code Change
#29013 opened May 27, 2025 by copybara-service bot Loading…
Bump hypothesis from 6.102.4 to 6.131.28 dependencies Pull requests that update a dependency file python Pull requests that update Python code
#29012 opened May 26, 2025 by dependabot bot Loading…
@jakevdp
Bump fsspec from 2024.5.0 to 2025.5.1 dependencies Pull requests that update a dependency file python Pull requests that update Python code
#29011 opened May 26, 2025 by dependabot bot Loading…
[Mosaic GPU] Use warpgroup semantics for the ragged dot example kernel.
#29007 opened May 26, 2025 by copybara-service bot Loading…
[Mosaic GPU] Resolve different tile transforms using the largest common divisor.
#29006 opened May 26, 2025 by copybara-service bot Loading…
Tighten numpy.linalg.qr annotation
#28997 opened May 24, 2025 by NeilGirdhar Loading…
@dfm
4
[scan] don't hoist loop-invariant computations in scan, just forward kokoro:force-run pull ready Ready for copybara import and testing
#28993 opened May 23, 2025 by mattjj • Draft
@mattjj
2
[JAX][DOC] Add optimizer state offloading doc
#28988 opened May 23, 2025 by zhenying-liu Loading…
3
[Mosaic:TPU] 16-bit tpu.iota
#28987 opened May 23, 2025 by copybara-service bot Loading…
Make CI job names to be shorter
#28986 opened May 23, 2025 by copybara-service bot Loading…
@nitins17
Apply extensive input to extensive output forwarding in scan. kokoro:force-run pull ready Ready for copybara import and testing
#28985 opened May 23, 2025 by dfm Loading…
@dfm
12
Preload the nvJitLink library installed by pip if available
#28977 opened May 23, 2025 by dfm Loading…
1
@dfm
Use the builtin isinstance with MLIR types pull ready Ready for copybara import and testing
#28973 opened May 23, 2025 by superbobry Loading…
[TSAN CI] Make jax build/test step fail if missing deps wheels kokoro:force-run pull ready Ready for copybara import and testing
#28972 opened May 23, 2025 by vfdev-5 Loading…
ProTip! Add no:assignee to see everything that’s not assigned.