Skip to content

Pull requests: jax-ml/jax

New pull request New
446 Open 17,462 Closed
446 Open 17,462 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

Warn the user if transparent huge pages aren't enabled.
#26978 opened Mar 6, 2025 by copybara-service bot Loading…
@mwhittaker
Prevent tracer leaks in scipy.special.expn
#26977 opened Mar 6, 2025 by jakevdp Loading…
1
@jakevdp
[Pallas/Fuser] Add debug option to fuser.fuse that prints out jaxpr
#26976 opened Mar 6, 2025 by copybara-service bot Loading…
@sharadmv
Update install docs for AMDGPU
#26975 opened Mar 6, 2025 by mrodden Loading…
jax.random.uniform: add resolution='high' argument
#26974 opened Mar 6, 2025 by jakevdp Loading…
@jakevdp
Expose profiler_data submodule from XLA to Jaxlib.
#26970 opened Mar 6, 2025 by copybara-service bot Loading…
Define lax.ragged_dot_general and express lax.ragged_dot in terms of it.
#26969 opened Mar 6, 2025 by copybara-service bot Loading…
Fix: Add type hints for pmap in_axes
#26966 opened Mar 6, 2025 by sarvagyad2604 Loading…
3
AutoPGLE: force-disable graphs less
#26964 opened Mar 6, 2025 by olupton Loading…
[Mosaic TPU][Pallas] Add pltpu.reciprocal_approx
#26960 opened Mar 6, 2025 by copybara-service bot Loading…
add autodiff rules for jax.lax.ragged_all_to_all collective
#26959 opened Mar 6, 2025 by copybara-service bot Loading…
@mattjj
[Mosaic TPU] Support fp8 upcast to f32
#26957 opened Mar 6, 2025 by copybara-service bot Loading…
Redacted
#26955 opened Mar 5, 2025 by copybara-service bot Loading…
add pmap axes hints pull ready Ready for copybara import and testing
#26950 opened Mar 5, 2025 by lockwo Loading…
10
Delete PjRtClient.Defragment.
#26948 opened Mar 5, 2025 by copybara-service bot Loading…
Replace pypi wheel dependencies with py_import.
#26945 opened Mar 5, 2025 by copybara-service bot Loading…
Support NVFP4 in scaled_matmul and scaled_dot_general
#26944 opened Mar 5, 2025 by wenscarl • Draft
Add NSA Attention implementation
#26943 opened Mar 5, 2025 by vishwamartur Loading…
1
Update remat config name in docs
#26929 opened Mar 5, 2025 by danijar Loading…
[Mosaic] Parsing TPU version with regex.
#26925 opened Mar 5, 2025 by copybara-service bot Loading…
roofline: Support computing flops for unary ops.
#26923 opened Mar 5, 2025 by copybara-service bot Loading…
Add accuracy field to unary ops
#26922 opened Mar 4, 2025 by copybara-service bot Loading…
Retry connection errors to combat flakiness in google metadata server during tpu initialization (reopens #24432)
#26919 opened Mar 4, 2025 by dlwh Loading…
[JAX] Add missing preset for X9 dot optimization on BF16/BF16 -> F32.
#26917 opened Mar 4, 2025 by copybara-service bot Loading…
[jax.distributed] Allow overriding XLA boot_id
#26906 opened Mar 4, 2025 by gspschmid Loading…
ProTip! Exclude everything labeled bug with -label:bug.