Skip to content

Pull requests: jax-ml/jax

New pull request New
Clear current search query, filters, and sorts
279 Open 8,203 Closed
279 Open 8,203 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

[pallas:mgpu] warp specialized pipeline with special delay_release.
#31483 opened Sep 2, 2025 by copybara-service bot Loading…
Canonicalize to vclamps (for vllm fp8 quant, maybe other uses)
#31452 opened Aug 31, 2025 by copybara-service bot Loading…
docs: Add note about computation perturbation in debug.print/breakpoint
#26682 opened Feb 22, 2025 by alhridoy Loading…
1
8
Clarify wait_recv semantics
#30000 opened Jul 7, 2025 by tengyifei Loading…
3
garbage in prettyprint
#32042 opened Sep 23, 2025 by mattjj • Draft
Add remote_bytes_transferred to Pallas CostEstimate dataclass.
#31760 opened Sep 11, 2025 by copybara-service bot Loading…
[Mosaic:TPU][NFC] Simplify computeSqueezeDimsChecked
#31652 opened Sep 9, 2025 by copybara-service bot Loading…
Set default odeint tol appropriately for float32 cla: yes
#3444 opened Jun 15, 2020 by shoyer Loading…
1
2
Specify init step size for odeint cla: yes
#2604 opened Apr 4, 2020 by jacobjinkelly • Draft
6
Set the default value of jax_ragged_dot_use_ragged_dot_instruction to True.
#32389 opened Oct 6, 2025 by copybara-service bot Loading…
feat(gh-13291): Add Pareto distribution functions: cdf, logcdf, logsf, sf, ppf, logppf, and corresponding tests
#26556 opened Feb 15, 2025 by Qazalbash Loading…
[pallas:mgpu] emit_pipeline_emit_pipeline to support different number of stages for each input.
#31543 opened Sep 4, 2025 by copybara-service bot Loading…
Integrate Triton up to [7decd434](https://github.com/openai/triton/commits/7decd434a2dbf76e74ea3c2a5ec877326b83f4fa)
#31554 opened Sep 4, 2025 by copybara-service bot Loading…
Add sample_image.
#30872 opened Aug 9, 2025 by carlosgmartin Loading…
4
#sdy Remove MHLO shardings from round-trip export
#30091 opened Jul 9, 2025 by copybara-service bot Loading…
Support building with Bzlmod
#32029 opened Sep 23, 2025 by meteorcloudy • Draft
2
Checkify: instrumentation context manager
#15389 opened Apr 4, 2023 by LenaMartens • Draft
11
HMM EM example cla: yes
#2405 opened Mar 12, 2020 by mattjj • Draft
FFT Convolution cla: yes
#6343 opened Apr 5, 2021 by peterroelants Loading…
12
Prefix PRNG Part 2: implement 1D prefix PRNG cla: yes
#2181 opened Feb 6, 2020 by j-towns Loading…
17
Add scan to checkify docs and fix rendering of ErrorCategories in docs.
#12079 opened Aug 24, 2022 by LenaMartens Loading…
1
4
Run scalar integer index normalization at trace time
#10683 opened May 12, 2022 by lgeiger Loading…
1
3
Support mask of broadcast, consts (arange, zeros, ...) cla: yes
#3472 opened Jun 17, 2020 by juliuskunze Loading…
Added doc: more cases in troubleshooting profiling
#16598 opened Jun 30, 2023 by OhadRubin Loading…
avals with names cla: yes
#4397 opened Sep 24, 2020 by jekbradbury • Draft
1
ProTip! Mix and match filters to narrow down what you’re looking for.