Skip to content

Pull requests: jax-ml/jax

New pull request New
433 Open 17,259 Closed
433 Open 17,259 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

Use uv instead of pip for installing Python packages
#26705 opened Feb 24, 2025 by copybara-service bot Loading…
@nitins17
[jax:custom_partitioning] Propagate static arguments to sharding_rule callback.
#26704 opened Feb 24, 2025 by copybara-service bot Loading…
Use LAPACK trsm kernel even for batched solves.
#26703 opened Feb 24, 2025 by copybara-service bot Loading…
@dfm
Add tests for lax.linalg.svd algorithm specification.
#26702 opened Feb 24, 2025 by copybara-service bot Loading…
@dfm
Cache the creating of sharding on ShapedArray since it's expensive to do it a lot of times
#26701 opened Feb 24, 2025 by copybara-service bot Loading…
@yashk2810
Use the same name for aliased Vars when pretty-printing Jaxprs. pull ready Ready for copybara import and testing
#26697 opened Feb 24, 2025 by gnecula Loading…
@gnecula
[Pallas] Add support for setting operand as HBM memory.
#26694 opened Feb 24, 2025 by copybara-service bot Loading…
fix member access to packed CUDA struct
#26691 opened Feb 23, 2025 by h-vetinari Loading…
@apaszke
Replace setup.py with pyproject.toml
#26683 opened Feb 23, 2025 by NeilGirdhar Loading…
@hawkinsp
docs: Add note about computation perturbation in debug.print/breakpoint
#26682 opened Feb 22, 2025 by alhridoy Loading…
1
7
Allow query and keys that aren’t multiples of 128
#26676 opened Feb 22, 2025 by Rifur13 Loading…
1
@sharadmv
[Mosaic-GPU] [3/3] Add support for communication primitives in MGPU lowering
#26675 opened Feb 21, 2025 by nvcastet Loading…
[Mosaic-GPU] [2/3] Add NVSHMEM support to Mosaic-GPU custom call
#26674 opened Feb 21, 2025 by nvcastet Loading…
[Pallas] [1/3] Move communication primitives from mosaic to core
#26673 opened Feb 21, 2025 by nvcastet Loading…
Enable :debug_info_test/test_vjp_of_jit. CI Connection Halt - Always Flag every job that has its workflow file setup for halting to wait at the designated point CI Connection Halt - On Error Flag every job that has its workflow file setup for halting to wait using trap CI Connection Halt - On Retry Flag every job that has its workflow file setup for halting to halt if the workflow is retried
#26671 opened Feb 21, 2025 by copybara-service bot Loading…
@danielsuo
Redefine is_fully_addressable in shardings to support zero local devices for McJAX.
#26670 opened Feb 21, 2025 by copybara-service bot Loading…
@emilyfertig
Sharp bits: add note on subnormal flush-to-zero documentation pull ready Ready for copybara import and testing
#26668 opened Feb 21, 2025 by jakevdp Loading…
1
@jakevdp
[mgpu] Debug print arrays.
#26659 opened Feb 21, 2025 by copybara-service bot Loading…
[Mosaic GPU] Use Index instead of I32 for the indices of async_load and async_store.
#26655 opened Feb 21, 2025 by copybara-service bot Loading…
Internal change for visibility
#26654 opened Feb 21, 2025 by copybara-service bot Loading…
[Pallas] Add option for async DMAs in the new TPU interpret mode
#26650 opened Feb 20, 2025 by copybara-service bot Loading…
@jburnim
roofline: Add support for min_p, max_p, reduce_sum_p.
#26643 opened Feb 20, 2025 by copybara-service bot Loading…
jax.numpy ndim/shape/size: deprecate non-array input pull ready Ready for copybara import and testing
#26641 opened Feb 20, 2025 by jakevdp Loading…
1
@jakevdp
[Mosaic GPU] Wire up the slice_lengths and indices operands in lowering of the MLIR dialect.
#26633 opened Feb 20, 2025 by copybara-service bot Loading…
Remove fmha rewriter dependency in cuDNN SDPA api
#26627 opened Feb 20, 2025 by Cjkkkk Loading…
ProTip! Exclude everything labeled bug with -label:bug.