Skip to content

Pull requests: jax-ml/jax

New pull request New
420 Open 16,807 Closed
420 Open 16,807 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

Optimize JaxprEqnContext context manager.
#26109 opened Jan 26, 2025 by copybara-service bot Loading…
@hawkinsp
enable setting flags via env vars like JAX_CPU_COLLECTIVES_IMPLEMENTATION
#26103 opened Jan 26, 2025 by mattjj Loading…
[better_errors] Refactor more uses of pe.tracing_debug_info (part 3) pull ready Ready for copybara import and testing
#26100 opened Jan 25, 2025 by gnecula Loading…
@gnecula
[better_errors] Refactor more uses of pe.tracing_debug_info (part 2) pull ready Ready for copybara import and testing
#26099 opened Jan 25, 2025 by gnecula Loading…
@gnecula
[better_errors] Refactor more uses of pe.tracing_debug_info (part 1) pull ready Ready for copybara import and testing
#26097 opened Jan 25, 2025 by gnecula Loading…
@gnecula
[Mosaic TPU] Support dynamic gather along axis 0 or 1 for 32-bit vreg-sized vector.
#26095 opened Jan 25, 2025 by copybara-service bot Loading…
Reject invalid None in jax.NamedSharding(spec=None).
#26090 opened Jan 24, 2025 by copybara-service bot Loading…
@pschuh
Update JVP rule for abs to fix behavior for complex infinite inputs pull ready Ready for copybara import and testing
#26086 opened Jan 24, 2025 by dfm Loading…
@dfm
11
[pallas:triton] The lowering now uses PTX instead of Triton IR
#26081 opened Jan 24, 2025 by copybara-service bot Loading…
@superbobry
[Mosaic GPU] Fix error message to make it clearer.
#26080 opened Jan 24, 2025 by copybara-service bot Loading…
[better_errors] Add debug info to more Jaxprs and Wrappedfun pull ready Ready for copybara import and testing
#26078 opened Jan 24, 2025 by gnecula • Draft
@gnecula
Add get/set methods to TransformedRef
#26075 opened Jan 24, 2025 by copybara-service bot Loading…
@sharadmv
Internal Change. DO NOT SUBMIT: Remove this line after the baseline is
#26071 opened Jan 23, 2025 by copybara-service bot Loading…
JEP: pytree preservation of dict key order
#26069 opened Jan 23, 2025 by jakevdp Loading…
5
Update --config=cuda to add CUDA dependencies both for bazel build and bazel test phases.
#26068 opened Jan 23, 2025 by copybara-service bot Loading…
Fix import for Windows platforms
#26066 opened Jan 23, 2025 by copybara-service bot Loading…
@Rifur13
[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…
@gnecula
[Mosaic GPU] Add support for tiled loads/stores with sub-byte types
#26057 opened Jan 23, 2025 by copybara-service bot Loading…
@apaszke
Initialize BlockSizes from attention.BlockSizes only on non-Windows platforms
#26052 opened Jan 23, 2025 by copybara-service bot Loading…
@nitins17
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…
@jakevdp
[Mosaic GPU] Use a single instance of the single_thread_predicate in the MLIR dialect lowering.
#26041 opened Jan 22, 2025 by copybara-service bot Loading…
[Mosaic GPU] Add specialized support for some int4 -> bfloat16 casts
#26039 opened Jan 22, 2025 by copybara-service bot Loading…
@apaszke
Remove linear_util.cache pull ready Ready for copybara import and testing
#26038 opened Jan 22, 2025 by dougalm Loading…
[Mosaic GPU] Handle the swizzle attribute in the lowering of async_store and async_load
#26037 opened Jan 22, 2025 by copybara-service bot Loading…
[mosaic_gpu] Allow tiled array instead of wgmma.
#26033 opened Jan 22, 2025 by copybara-service bot Loading…
ProTip! Find all pull requests that aren't related to any open issues with -linked:issue.