Skip to content

Pull requests: jax-ml/jax

New pull request New
412 Open 15,894 Closed
412 Open 15,894 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

Update jax.scipy.special.gamma and gammasgn to return NaN for negativ…
#24945 opened Nov 18, 2024 by hawkinsp Loading…
[pallas:mosaic_gpu] Cache IR constants in the lowering
#24944 opened Nov 18, 2024 by copybara-service bot Loading…
@superbobry
[mosaic_gpu] Fixed FragmentedArray comparisons with literals
#24943 opened Nov 18, 2024 by copybara-service bot Loading…
@superbobry
fix(docs): typos in macro name in FFI tutorial
#24942 opened Nov 18, 2024 by jeertmans Loading…
[Mosaic][Easy] - Wire up kernel names to MLIR dump
#24936 opened Nov 17, 2024 by copybara-service bot Loading…
Use a direct StableHLO lowering for pow. pull ready Ready for copybara import and testing
#24933 opened Nov 16, 2024 by hawkinsp Loading…
Return a ndarray in shape_as_value if the shape is known to be constant. pull ready Ready for copybara import and testing
#24932 opened Nov 16, 2024 by hawkinsp Loading…
Add an allow_negative_indices option to lax.dynamic_slice and lax.dynamic_update_slice.
#24931 opened Nov 16, 2024 by copybara-service bot Loading…
@hawkinsp
[Pallas] Support unblocked indexing mode in emit_pipeline.
#24928 opened Nov 16, 2024 by copybara-service bot Loading…
Update jax.tree module docs
#24925 opened Nov 15, 2024 by 8bitmp3 • Draft
Update jax.grad docs
#24924 opened Nov 15, 2024 by 8bitmp3 • Draft
Update jax.random docs
#24923 opened Nov 15, 2024 by 8bitmp3 • Draft
Update jax.jit docs
#24921 opened Nov 15, 2024 by 8bitmp3 • Draft
5
Update jax.numpy docs
#24920 opened Nov 15, 2024 by 8bitmp3 • Draft
[Mosaic TPU] Fold sublane offset to indices when storing to untiled ref.
#24919 opened Nov 15, 2024 by copybara-service bot Loading…
DOC: update main landing page style
#24916 opened Nov 15, 2024 by jakevdp Loading…
@jakevdp
Set __module__ attribute for objects in jax.numpy pull ready Ready for copybara import and testing
#24912 opened Nov 15, 2024 by jakevdp Loading…
@jakevdp
1
Add JAX_COMPILATION_CACHE_EXPECT_PGLE option
#24910 opened Nov 15, 2024 by olupton Loading…
C++ tree with path API
#24906 opened Nov 14, 2024 by copybara-service bot Loading…
@IvyZX
Make logaddexp and logaddexp2 into ufuncs pull ready Ready for copybara import and testing
#24903 opened Nov 14, 2024 by jakevdp Loading…
@jakevdp
2
[Pallas TPU] Support masked store
#24902 opened Nov 14, 2024 by copybara-service bot Loading…
[pallas] Minor simplifications to Pallas interpreter.
#24900 opened Nov 14, 2024 by copybara-service bot Loading…
@chr1sj0nes
Added CI job with TSAN and free-threading
#24898 opened Nov 14, 2024 by vfdev-5 • Draft
[Mosaic TPU] Add general tpu.vector_store and support masked store.
#24892 opened Nov 14, 2024 by copybara-service bot Loading…
[Mosaic TPU] Support relayout for mask vector
#24891 opened Nov 13, 2024 by copybara-service bot Loading…
ProTip! Add no:assignee to see everything that’s not assigned.