Skip to content

Pull requests: jax-ml/jax

New pull request New
455 Open 16,448 Closed
455 Open 16,448 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

Add SMEM as a supported Pallas output memory space.
#25689 opened Dec 27, 2024 by copybara-service bot Loading…
Add jax.random.multinomial.
#25688 opened Dec 27, 2024 by carlosgmartin Loading…
1
2
Fix formatting in docs for transposing pytrees pull ready Ready for copybara import and testing
#25686 opened Dec 26, 2024 by Mikcl Loading…
@dfm
2
[Mosaic:TPU][NFC] Clean up unused variable
#25680 opened Dec 24, 2024 by copybara-service bot Loading…
Deprecate scipy.special.lpmn & lpmn_values pull ready Ready for copybara import and testing
#25675 opened Dec 23, 2024 by jakevdp Loading…
@jakevdp
Expand test case to repro a crash.
#25670 opened Dec 23, 2024 by jburnim Loading…
Tensorboard profiling plugin nightly instructions pull ready Ready for copybara import and testing
#25661 opened Dec 23, 2024 by rdyro Loading…
Add support for axis_name and axis_index_groups to lax.ragged_all_to_all
#25660 opened Dec 23, 2024 by copybara-service bot Loading…
Simplify implementation of random.orthogonal.
#25655 opened Dec 22, 2024 by carlosgmartin Loading…
Fix attention backwards pass
#25647 opened Dec 20, 2024 by Rifur13 Loading…
[Mosaic TPU] Validate inserted layout in relayout-insertion pass.
#25645 opened Dec 20, 2024 by copybara-service bot Loading…
Fix backwards attention test for larger head dims
#25642 opened Dec 20, 2024 by Rifur13 Loading…
[pallas_mgpu] For loops can have accumulators for carries.
#25640 opened Dec 20, 2024 by copybara-service bot Loading…
[pallas/pallas_mgpu] Discharging run_scoped should not be discharging the intermediates
#25639 opened Dec 20, 2024 by copybara-service bot Loading…
[Pallas] Add a test to ensure consistent rounding during float-to-int casts
#25637 opened Dec 20, 2024 by copybara-service bot Loading…
@apaszke
[Mosaic TPU] Add support for second minor broadcasts with packed types
#25636 opened Dec 20, 2024 by copybara-service bot Loading…
@apaszke
Move jax.extend.ffi module to top level jax.ffi pull ready Ready for copybara import and testing
#25633 opened Dec 20, 2024 by dfm Loading…
@dfm
3
[Mosaic:TPU] Less urgent fixes to trunc infer rule after cl/708011538
#25627 opened Dec 20, 2024 by copybara-service bot Loading…
Switch to a new thread-safe utility for catching warnings.
#25626 opened Dec 20, 2024 by hawkinsp Loading…
5
Remove dead codepaths now that MemorySpaceDescription works in OSS
#25621 opened Dec 19, 2024 by copybara-service bot Loading…
@yashk2810
[Mosaic GPU] Clean up remnants of the early C++ dialect implementation.
#25610 opened Dec 19, 2024 by copybara-service bot Loading…
@bchetioui
[Mosaic TPU] Add support for true divide in bf16 on TPUv6
#25608 opened Dec 19, 2024 by copybara-service bot Loading…
@apaszke
[Mosaic TPU] Be much more aggressive in inferring large 2nd minor layouts for 16-bit types on v6
#25605 opened Dec 19, 2024 by copybara-service bot Loading…
@apaszke
[Mosaic] NFC: Pull out vreg related functions to util.
#25600 opened Dec 19, 2024 by copybara-service bot Loading…
Add examples for jax.checkpoint activation offloading policies pull ready Ready for copybara import and testing
#25594 opened Dec 18, 2024 by zhenying-liu Loading…
2
ProTip! Adding no:label will show everything without a label.