Skip to content

Pull requests: jax-ml/jax

New pull request New
508 Open 19,094 Closed
508 Open 19,094 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

Tighten numpy.linalg.qr annotation
#28997 opened May 24, 2025 by NeilGirdhar Loading…
1
Automated Code Change
#28995 opened May 24, 2025 by copybara-service bot Loading…
Block until ready for PGLE test
#28994 opened May 24, 2025 by copybara-service bot Loading…
[scan] don't hoist loop-invariant computations in scan, just forward pull ready Ready for copybara import and testing
#28993 opened May 23, 2025 by mattjj • Draft
@mattjj
[JAX][DOC] Add optimizer state offloading doc
#28988 opened May 23, 2025 by zhenying-liu Loading…
1
[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.
#28985 opened May 23, 2025 by dfm Loading…
@dfm
5
Move jax/_src/custom_derivatives.py to its own BUILD rule
#28982 opened May 23, 2025 by copybara-service bot Loading…
@jakevdp
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…
[Pallas:MGPU] Support remote async copies and use them in the collective matmul
#28970 opened May 23, 2025 by copybara-service bot Loading…
@apaszke
Parametrize build system on CUDA major version
#28968 opened May 23, 2025 by olupton Loading…
[TSAN CI] Try NO_CYTHON_COMPILE=true when installing Cython
#28967 opened May 23, 2025 by vfdev-5 • Draft
Explicitly make out a jax.Array before converting it in lax_numpy.array.
#28966 opened May 23, 2025 by copybara-service bot Loading…
@bchetioui
Expose Recompute in partial_eval module.
#28956 opened May 22, 2025 by copybara-service bot Loading…
Avoid doing DCE of effectful ops and reordering in partial eval. pull ready Ready for copybara import and testing
#28955 opened May 22, 2025 by dougalm Loading…
1
[Mosaic GPU] Add non-collective blackwell matmul example CI Optional GPU Presubmit Label to flag PR to run additional GPU testing not in standard presubmits
#28954 opened May 22, 2025 by copybara-service bot Loading…
@justinjfu
Initial commit for attaching XLA metadata to individual HLO operations via 'jax.debug.attach_metadata(...)'
#28953 opened May 22, 2025 by copybara-service bot Loading…
[Mosaic:TPU] Byte-granularity dynamic gathers
#28952 opened May 22, 2025 by copybara-service bot Loading…
Add extra cuBLAS/cuDNN version checks
#28931 opened May 22, 2025 by olupton Loading…
[draft][Mosaic:TPU][Relayout] Implicit minor -> implicit 2nd minor for 16-bit and 8-bit (TPUv6+)
#28924 opened May 22, 2025 by copybara-service bot Loading…
Add a linearization rule for cond_p.
#28919 opened May 22, 2025 by dfm Loading…
@dfm
[jax2tf] Refine the disabling of jax2tf_test, for versions <= 2.19.1 kokoro:force-run pull ready Ready for copybara import and testing
#28916 opened May 21, 2025 by gnecula Loading…
@gnecula
ProTip! Add no:assignee to see everything that’s not assigned.