Skip to content

Pull requests: jax-ml/jax

New pull request New
469 Open 18,551 Closed
469 Open 18,551 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

Fix overloaded type signature for jax.numpy.where.
#28314 opened Apr 27, 2025 by carlosgmartin Loading…
Make primary sidebar open by default and toggleable.
#28313 opened Apr 27, 2025 by zhangqiaorjc Loading…
1
Automated Code Change
#28310 opened Apr 26, 2025 by copybara-service bot Loading…
[Mosaic TPU] Fold reshape (..., M, N, 128) -> (..., M, N * 128) with any tiling and dtype to load.
#28309 opened Apr 26, 2025 by copybara-service bot Loading…
#sdy raise error if module was exported with Shardy disabled but loaded & called with Shardy enabled.
#28307 opened Apr 26, 2025 by copybara-service bot Loading…
[Mosaic] Add none memory space. Add a test for copying a (1,) shaped array to SMEM.
#28306 opened Apr 26, 2025 by copybara-service bot Loading…
Uses include to reference xla packages in jaxlib xla BUILD
#28304 opened Apr 25, 2025 by yliu120 Loading…
[Mosaic] Remove Python pipeline.
#28303 opened Apr 25, 2025 by copybara-service bot Loading…
[Pallas][Mosaic GPU] Add support for blackwell mma instruction.
#28302 opened Apr 25, 2025 by copybara-service bot Loading…
@justinjfu
[JAX] Add a python 3.14 requirement lock file and update WORKSPACE
#28301 opened Apr 25, 2025 by copybara-service bot Loading…
Add is_leaf_with_path predicate.
#28300 opened Apr 25, 2025 by copybara-service bot Loading…
1
@IvyZX
Add support for output and input memory space colors in tpu custom calls via CustomCallConfig.
#28290 opened Apr 25, 2025 by copybara-service bot Loading…
Declare tpu.vector_load, mirroring tpu.vector_store.
#28289 opened Apr 25, 2025 by copybara-service bot Loading…
Ensure __jax_array__ works properly with JIT disabled. pull ready Ready for copybara import and testing
#28288 opened Apr 25, 2025 by jakevdp Loading…
@jakevdp
Add debug.capture() callback wrapper for argument-free debug printing
#28284 opened Apr 25, 2025 by kcdodd Loading…
1
Add retry to certain K8s-related calls to overcome transient failure
#28283 opened Apr 25, 2025 by yhtang Loading…
Add argnums param to fwd_and_bwd
#28279 opened Apr 25, 2025 by gspschmid Loading…
Clarify that jnp.clip gives preference to maximum values where bounds are incongruent
#28275 opened Apr 25, 2025 by johannahaffner Loading…
@jakevdp
3
Trying to lower the degree of parallelism CI Optional GPU Presubmit Label to flag PR to run additional GPU testing not in standard presubmits
#28273 opened Apr 25, 2025 by copybara-service bot Loading…
@apaszke
[Mosaic GPU] Make sure the tests all pass on B200 (esp. in our CI)
#28272 opened Apr 25, 2025 by copybara-service bot Loading…
@apaszke
[Mosaic GPU] Add layout inference and lowering for vector.MultiDimReduce
#28271 opened Apr 25, 2025 by copybara-service bot Loading…
[Mosaic GPU] Handle WGMMA_ROW_LAYOUT in vector_store lowering and inference.
#28266 opened Apr 25, 2025 by copybara-service bot Loading…
[DO NOT MERGE] Relaxing B200 presubmit conditions
#28263 opened Apr 25, 2025 by copybara-service bot Loading…
@apaszke
[JAX:SPARSE] Implement CSR sparse kernel
#28261 opened Apr 25, 2025 by copybara-service bot Loading…
[pallas/mosaic] emulate fp32<->ui32 casts in terms of fp32<->si32 casts
#28257 opened Apr 25, 2025 by copybara-service bot Loading…
@jakevdp
ProTip! Type g i on any issue or pull request to go back to the issue listing page.