Skip to content

Issues: jax-ml/jax

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

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
Assignee
Filter by who’s assigned
Sort

Issues list

[better_errors] Add debug info to more Jaxprs pull ready Ready for copybara import and testing
#26078 opened Jan 24, 2025 by gnecula Loading…
[export] Fix misuse of NamedSharding in export tests pull ready Ready for copybara import and testing
#26076 opened Jan 24, 2025 by gnecula Loading…
Add jax.test_util to public API docs pull ready Ready for copybara import and testing
#26072 opened Jan 24, 2025 by shoyer Loading…
[better_errors] Port the Pallas debug info mechanisms to the new JAX debugging info pull ready Ready for copybara import and testing
#26058 opened Jan 23, 2025 by gnecula Loading…
increase sdpa dbias unit test error tolerance pull ready Ready for copybara import and testing
#26046 opened Jan 22, 2025 by Cjkkkk Loading…
Implement lax.full_like in terms of a primitive pull ready Ready for copybara import and testing
#26045 opened Jan 22, 2025 by jakevdp Loading…
Remove linear_util.cache pull ready Ready for copybara import and testing
#26038 opened Jan 22, 2025 by dougalm Loading…
[better_errors] Consolidate TracingDebugInfo and JaxprDebugInfo pull ready Ready for copybara import and testing
#26001 opened Jan 21, 2025 by gnecula Loading…
Clarify type hints for wrapped functions pull ready Ready for copybara import and testing
#25994 opened Jan 21, 2025 by ncoghlan Loading…
Add DCE rules for custom_jvp and custom_vjp pull ready Ready for copybara import and testing
#25963 opened Jan 17, 2025 by dfm Loading…
[shard_map] fix jax.debug.print/callback under partial auto pull ready Ready for copybara import and testing
#25887 opened Jan 14, 2025 by mattjj Loading…
Adding GPU paged attention kernel pull ready Ready for copybara import and testing
#25839 opened Jan 10, 2025 by Rifur13 Loading…
Add cuDNN flash attention sequence packing pull ready Ready for copybara import and testing
#25812 opened Jan 9, 2025 by Cjkkkk Loading…
Add jax.random.multinomial. kokoro:force-run pull ready Ready for copybara import and testing
#25688 opened Dec 27, 2024 by carlosgmartin Loading…
Fix debug_nans regressions. pull ready Ready for copybara import and testing
#25519 opened Dec 16, 2024 by emilyfertig Loading…
Simplify implementation of nn.relu. pull ready Ready for copybara import and testing
#25331 opened Dec 8, 2024 by carlosgmartin Loading…
Support dynamic masks in splash attention kokoro:force-run pull ready Ready for copybara import and testing
#25213 opened Dec 2, 2024 by Rifur13 Loading…
improve jnp.mean / jnp.sum / ... error message for out-of-bounds axis index better_errors Improve the error reporting pull ready Ready for copybara import and testing
#25155 opened Nov 28, 2024 by mattjj Loading…
Start a new TPU interpret mode for Pallas pull ready Ready for copybara import and testing
#25097 opened Nov 25, 2024 by jburnim Loading…
Added stream annotation support via @compute_on('stream:#') decorator pull ready Ready for copybara import and testing
#25056 opened Nov 22, 2024 by chaserileyroberts Loading…
Implement SVD algorithm based on QR for CPU targets pull ready Ready for copybara import and testing
#25053 opened Nov 22, 2024 by JanLuca Loading…
allow custom_gradient to close over non-autodiff tracers pull ready Ready for copybara import and testing
#25017 opened Nov 20, 2024 by mattjj Loading…
[pallas] Expand shape polymorphism support for Pallas. pull ready Ready for copybara import and testing
#24960 opened Nov 18, 2024 by gnecula Loading…
Update jax.jit docs pull ready Ready for copybara import and testing
#24921 opened Nov 15, 2024 by 8bitmp3 Draft
ProTip! Follow long discussions with comments:>50.