Skip to content

Pull requests: jax-ml/jax

New pull request New
533 Open 19,209 Closed
533 Open 19,209 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

Update workflow files to use new ml-build containers.
#29171 opened Jun 2, 2025 by copybara-service bot Loading…
[cleanup] inline uses of NumpyComplexWarning pull ready Ready for copybara import and testing
#29170 opened Jun 2, 2025 by jakevdp Loading…
@jakevdp
always compile Pallas calls, enabling pallas_call under disable_jit kokoro:force-run pull ready Ready for copybara import and testing
#29168 opened Jun 2, 2025 by froystig Loading…
@froystig
2
Add sharding rule for cudnn attention kernels, whose doc is at https://docs.nvidia.com/deeplearning/cudnn/frontend/latest/operations/Attention.html#attention.
#29163 opened Jun 2, 2025 by copybara-service bot Loading…
#sdy Have JAX export compat tests also run on Shardy.
#29161 opened Jun 2, 2025 by copybara-service bot Loading…
Try to run tsan.yaml on linux-arm64-c4a-64
#29160 opened Jun 2, 2025 by vfdev-5 • Draft
[jax] Always mask padded region if query_seq_lengths is passed to dot_product_attention.
#29158 opened Jun 2, 2025 by copybara-service bot Loading…
@chr1sj0nes
Fix int type promotion in multihost_utils.broadcast_one_to_all.
#29157 opened Jun 2, 2025 by heiner Loading…
2
Add initial devcontainer configuration for Alpine environment
#29154 opened Jun 1, 2025 by agluhic Loading…
3
Simplify jnp.isclose
#29153 opened Jun 1, 2025 by soraros Loading…
4
[xla:cpu] Deprecate API_VERSION_STATUS_RETURNING custom calls
#29150 opened Jun 1, 2025 by copybara-service bot Loading…
@ezhulenev
Implementing PReLu
#29147 opened May 31, 2025 by MythicArrow Loading…
@jakevdp
1
Fix typos discovered by codespell
#29146 opened May 31, 2025 by cclauss Loading…
@jakevdp
5
Adding tests and improving jnp.ufunc support kokoro:force-run pull ready Ready for copybara import and testing
#29144 opened May 30, 2025 by DanisNone Loading…
12
Bug fix in PjRt Cuda plugin to correct mismatched device ordinals in stream executor/device default.
#29142 opened May 30, 2025 by copybara-service bot Loading…
@emilyfertig
[Mosaic GPU] Error when causal masking is used on cuda versions known to result in a ptxas miscompilation (between 12.8.0 and 12.9.1).
#29141 opened May 30, 2025 by copybara-service bot Loading…
@Rifur13
Support/test PjRt C API executables with no addressable devices.
#29138 opened May 30, 2025 by copybara-service bot Loading…
@emilyfertig
add psend and precv to jax/lax/parallel
#29135 opened May 30, 2025 by rosiezou Loading…
17
Fix a rare numerical flake in svd_test seen on TPU v6e. pull ready Ready for copybara import and testing
#29132 opened May 30, 2025 by hawkinsp Loading…
[XProf] Change tensorboard-plugin-profile to new xprof package documentation
#29129 opened May 30, 2025 by Matt-Hurd Loading…
@ZacCranko
[JAX] Remove the redundant pjit BUILD target.
#29126 opened May 30, 2025 by copybara-service bot Loading…
@jakevdp
Change all us-central1-docker.pkg.dev/tensorflow-sigs/tensorflow/ml-build-arm64 container to us-docker.pkg.dev/ml-oss-artifacts-published/ml-public-container/ml-build-arm64.
#29124 opened May 30, 2025 by copybara-service bot Loading…
[pallas:mosaic_gpu] plgpu.nd_loop is now a decorator similar to pl.loop
#29123 opened May 30, 2025 by copybara-service bot Loading…
@superbobry
[Mosaic GPU] Convert all memrefs with transforms to unrealized casts and check them.
#29122 opened May 30, 2025 by copybara-service bot Loading…
[pallas:mosaic_gpu] Dropped the GPU prefix from GPUShapeDtypeStruct
#29119 opened May 30, 2025 by copybara-service bot Loading…
@superbobry
ProTip! Follow long discussions with comments:>50.