Skip to content

Pull requests: jax-ml/jax

New pull request New
457 Open 16,427 Closed
457 Open 16,427 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

Fix a typo in documentation for pinv function pull ready Ready for copybara import and testing
#25662 opened Dec 23, 2024 by liblaf Loading…
1
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…
Internal: use a single registry for abstractify APIs pull ready Ready for copybara import and testing
#25651 opened Dec 20, 2024 by jakevdp Loading…
@jakevdp
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
2
[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
[AutoPGLE] Explicitly disable command buffers when profiler is used.
#25607 opened Dec 19, 2024 by copybara-service bot Loading…
[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 GPU] Remove jaxlib version check for the Mosaic GPU dialect.
#25602 opened Dec 19, 2024 by copybara-service bot Loading…
@bchetioui
[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
In progress experimention. Add StringDType to JAX's supported types.
#25592 opened Dec 18, 2024 by copybara-service bot Loading…
#sdy enable pure callbacks and debug prints in JAX.
#25586 opened Dec 18, 2024 by copybara-service bot Loading…
ProTip! Find all pull requests that aren't related to any open issues with -linked:issue.