Skip to content

Pull requests: jax-ml/jax

New pull request New
556 Open 21,793 Closed
556 Open 21,793 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

Add support for device kwarg in astype, and add matching utility func
#21086 opened May 6, 2024 by Micky774 Loading…
28
add error for remat + rbg for #21046
#21049 opened May 3, 2024 by mattjj • Draft
4
make _ProfileState's reset point to self.
#21020 opened Apr 30, 2024 by sadikneipp Loading…
Use jax.numpy.diagonal to evaluate jax.numpy.trace.
#20970 opened Apr 28, 2024 by tillahoffmann Loading…
4
Add support for rounding integers with decimals<0 in jnp.round
#20829 opened Apr 19, 2024 by james77777778 Loading…
Minor formatting fix for jax.nn.one_hot docstring
#20828 opened Apr 19, 2024 by albanie Loading…
@jakevdp
2
Fix error message in paged_attention_kernel
#20791 opened Apr 16, 2024 by wonjoo-wj Loading…
Clarify that lax.top_k returns values and corresponding indices in descending order pull ready Ready for copybara import and testing
#20781 opened Apr 16, 2024 by Sohl-Dickstein Loading…
@hawkinsp
7
Fix to make jax.scipy.special.sph_harm to be close to scipy.special.sph_harm for multiple theta or phi
#20772 opened Apr 16, 2024 by rajasekharporeddy Loading…
1
9
Fix typo in docs
#20744 opened Apr 14, 2024 by yueqiw Loading…
@jakevdp
1
Add type information for pallas bitcast/roll pull ready Ready for copybara import and testing
#20576 opened Apr 3, 2024 by kkiningh Loading…
@sharadmv
5
Making JAX API Reference easier to find on documentation sidebar pull ready Ready for copybara import and testing
#20549 opened Apr 2, 2024 by ayaka14732 Loading…
1
@ayaka14732
1
removing enhanced barrier from cloud init file because it is getting on by default elsewhere
#20480 opened Mar 28, 2024 by sshahrokhi Loading…
1
if fori_loop statically takes zero trips, don't run body pull ready Ready for copybara import and testing
#20448 opened Mar 27, 2024 by mattjj Loading…
1
3
Flatten defaultdict with path in key-sorted order
#20423 opened Mar 25, 2024 by XuehaiPan Loading…
[Pallas/TPU] Add flash attention math tutorial
#20145 opened Mar 8, 2024 by sharadmv • Draft
1
Remove deprecated shims and update deprecated ruff config
#20116 opened Mar 7, 2024 by tttc3 Loading…
Auto-labels 'gemma' on 'gemma' issues/PRs.
#20113 opened Mar 7, 2024 by shmishra99 Loading…
try giving outside_call_p in/outfeed effects, tests fail
#19968 opened Feb 24, 2024 by mattjj • Draft
Handle jaxpr lists when pretty-printing key-value pairs
#19719 opened Feb 8, 2024 by Mandrenkov Loading…
1
Add jax.tree_util.tree_traverse pull ready Ready for copybara import and testing
#19695 opened Feb 7, 2024 by jakevdp Loading…
1
@jakevdp
10
Make DTypeLike runtime checkable
#19669 opened Feb 6, 2024 by pchanial Loading…
5
More type annotations of Pallas primitives pull ready Ready for copybara import and testing
#19608 opened Feb 1, 2024 by trishume Loading…
@sharadmv
4
jax.numpy: return tuple of array instead of list of array pull ready Ready for copybara import and testing
#19531 opened Jan 26, 2024 by jakevdp Loading…
@jakevdp
pallas: add argument zeroize_outputs to pallas_call()
#19399 opened Jan 17, 2024 by geggo Loading…
ProTip! What’s not been updated in a month: updated:<2025-08-29.