Skip to content

Pull requests: jax-ml/jax-tpu-embedding

New pull request New
7 Open 606 Closed
7 Open 606 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
Assigned to nobody Loading
Sort
Sort by
Newest Oldest Most commented Least commented Recently updated Least recently updated Best match
Most reactions
👍 👎 😄 🎉 😕 ❤️ 🚀 👀

Pull requests list

[JAX SC] Rename benchmarks for clarity
#614 opened Dec 5, 2025 by copybara-service bot Loading…
[JAX SC] Extract features slice per SC in parallel.
#612 opened Dec 5, 2025 by copybara-service bot Loading…
Add number of SparseCores for TPU v7 / Ironwood.
#610 opened Dec 4, 2025 by copybara-service bot Loading…
[JAX SC] Move benchmarks to benchmarks subfolder, tests to separate subfolder.
#602 opened Dec 2, 2025 by copybara-service bot Loading…
[JAX SC] Refactor input preprocessing to group data per SparseCore.
#590 opened Nov 28, 2025 by copybara-service bot Loading…
[JAX SC] required_buffer_size_per_sc is being undercounted when minibatching is enabled. The initial calculation does not account for the alignment padding added at the end of each minibatch.
#528 opened Oct 28, 2025 by copybara-service bot Loading…
Add presubmit script to jax-tpu-embedding/sparsecore CI Connection Halt - On Retry Flag every job that has its workflow file setup for halting to halt if the workflow is retried
#149 opened Jan 13, 2025 by copybara-service bot Loading…
ProTip! Type g i on any issue or pull request to go back to the issue listing page.