Skip to content

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

New pull request New
11 Open 584 Closed
11 Open 584 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] perf: Hoist can_drop_id && allow_id_dropping out of keys loop.
#596 opened Dec 1, 2025 by copybara-service bot Loading…
[JAX SC] perf: Optimize deduplication logic in sparse core.
#595 opened Dec 1, 2025 by copybara-service bot Loading…
[JAX SC] Update benchmarks to use enable_minibatching.
#594 opened Dec 1, 2025 by copybara-service bot Loading…
[JAX SC] Add combiner to sort benchmarks.
#593 opened Dec 1, 2025 by copybara-service bot Loading…
[JAX SC] Vectorize ragged input array generation (~16x speedup).
#592 opened Dec 1, 2025 by copybara-service bot Loading…
[JAX SC] Add id_count() to AbstractInputBatch for accurate pre-allocation.
#591 opened Nov 28, 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] Add support for loading EmbeddingSpecProto for benchmarking.
#588 opened Nov 25, 2025 by copybara-service bot Loading…
Optimize SparseCore input preprocessing by eliminating buffer copies.
#567 opened Nov 12, 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.