Skip to content

Pull requests: jax-ml/jax

New pull request New
462 Open 18,527 Closed
462 Open 18,527 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

[Mosaic GPU] Wire up the plugin that contains the MGPU custom call into Bazel tests
#28265 opened Apr 25, 2025 by copybara-service bot Loading…
@apaszke
[DO NOT MERGE] Relaxing B200 presubmit conditions
#28263 opened Apr 25, 2025 by copybara-service bot Loading…
@apaszke
[JAX:SPARSE] Implement CSR sparse kernel
#28261 opened Apr 25, 2025 by copybara-service bot Loading…
[Pallas] Fix index_map equality checks
#28259 opened Apr 25, 2025 by copybara-service bot Loading…
@sharadmv
Replace auto tracking with manual axis names tracking internally in shard_map.
#28258 opened Apr 25, 2025 by copybara-service bot Loading…
@yashk2810
[pallas/mosaic] emulate fp32<->ui32 casts in terms of fp32<->si32 casts
#28257 opened Apr 25, 2025 by copybara-service bot Loading…
@jakevdp
Add multiaccelerator H100 tests to optional GPU presubmit.
#28256 opened Apr 25, 2025 by copybara-service bot Loading…
[Mosaic] Support bf16 select and cmp <= TPUv4.
#28255 opened Apr 25, 2025 by copybara-service bot Loading…
1
Add conv_general_dilated sharding rule
#28253 opened Apr 24, 2025 by copybara-service bot Loading…
1
@yashk2810
[Mosaic] NFC: Refactor Mosaic canonicalization pass to use RewritePattern.
#28250 opened Apr 24, 2025 by copybara-service bot Loading…
Edit print_environment_info to print environment variables that start with JAX_. kokoro:force-run pull ready Ready for copybara import and testing
#28244 opened Apr 24, 2025 by carlosgmartin Loading…
@jakevdp
1
[pallas] Removed the platform-specific prefix from *CompilerParams and *MemorySpace
#28243 opened Apr 24, 2025 by copybara-service bot Loading…
@superbobry
[Mosaic GPU] Add a BroadcastInDim op in the mosaic mlir dialect.
#28237 opened Apr 24, 2025 by copybara-service bot Loading…
Update build.py to avoid duplication of bazel options both in bazel command and in .jax_configure.bazelrc.
#28236 opened Apr 24, 2025 by copybara-service bot Loading…
Re-land "Inline literals while tracing instead of in a separate pass".
#28233 opened Apr 24, 2025 by copybara-service bot Loading…
@dfm
Support pltpu.roll on sublanes when not all lanes are used.
#28228 opened Apr 24, 2025 by copybara-service bot Loading…
[Mosaic TPU] Allow simultaneous column and row shifts CI Optional GPU Presubmit Label to flag PR to run additional GPU testing not in standard presubmits
#28220 opened Apr 24, 2025 by copybara-service bot Loading…
@apaszke
[mosaic_gpu] Added debug_assert
#28215 opened Apr 24, 2025 by copybara-service bot Loading…
@superbobry
jax.numpy: move linspace & friends into array_creation kokoro:force-run pull ready Ready for copybara import and testing
#28209 opened Apr 23, 2025 by jakevdp Loading…
@jakevdp
style: type annotate flatten_util
#28201 opened Apr 23, 2025 by nstarman Loading…
[XLA] Replace XLA constructor from proto with factory method.
#28180 opened Apr 22, 2025 by copybara-service bot Loading…
Added additional stream annotation E2E tests kokoro:force-run pull ready Ready for copybara import and testing
#28170 opened Apr 22, 2025 by chaserileyroberts Loading…
7
Implement Gauss hypergeometric function 2F1
#28168 opened Apr 22, 2025 by mattbahr Loading…
37
Support mutability in the c++ dispatch.
#28167 opened Apr 22, 2025 by copybara-service bot Loading…
@pschuh
Warn on excessive captured constants documentation pull ready Ready for copybara import and testing
#28157 opened Apr 21, 2025 by ZacCranko Loading…
@ZacCranko
ProTip! What’s not been updated in a month: updated:<2025-03-25.