Skip to content

Pull requests: jax-ml/jax

New pull request New
512 Open 19,065 Closed
512 Open 19,065 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

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
Fix sharding rule in jax test
#28965 opened May 23, 2025 by copybara-service bot Loading…
Automated Code Change
#28963 opened May 23, 2025 by copybara-service bot Loading…
Transfer library: poison outstanding buffer fetches upon connection failure.
#28962 opened May 23, 2025 by copybara-service bot Loading…
@pschuh
Simplify attention VJP definition
#28960 opened May 23, 2025 by copybara-service bot Loading…
@rdyro
Add TODO to run TPU interpret mode tests in parallel.
#28958 opened May 22, 2025 by copybara-service bot Loading…
@jburnim
Move jax/_src/interpreters/batching.py into its own BUILD rule
#28957 opened May 22, 2025 by copybara-service bot Loading…
@jakevdp
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…
[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…
Update XlaCallModule so tests are compatible with DCE.
#28949 opened May 22, 2025 by copybara-service bot Loading…
Move _src/custom_transpose.py into its own BUILD rule
#28948 opened May 22, 2025 by copybara-service bot Loading…
@jakevdp
[Mosaic GPU] Add support for async copies to peer devices CI Optional GPU Presubmit Label to flag PR to run additional GPU testing not in standard presubmits
#28941 opened May 22, 2025 by copybara-service bot Loading…
@apaszke
[Mosaic GPU] Add checks for argument shapes and types
#28940 opened May 22, 2025 by copybara-service bot Loading…
@apaszke
Rename backend.compile to backend.compile_and_load.
#28936 opened May 22, 2025 by copybara-service bot Loading…
@danielsuo
[jaxlib] Add CompileOnlyPyClient to xla_client.
#28933 opened May 22, 2025 by copybara-service bot Loading…
@danielsuo
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
[Mosaic GPU] Add barrier transformation support to tcgen05_mma. CI Optional GPU Presubmit Label to flag PR to run additional GPU testing not in standard presubmits
#28917 opened May 22, 2025 by copybara-service bot Loading…
@justinjfu
[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! Exclude everything labeled bug with -label:bug.