Skip to content

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

New pull request New
10 Open 629 Closed
10 Open 629 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

Internal change.
#640 opened Dec 21, 2025 by copybara-service bot Loading…
Minor cleanup of the grad update op for Adam
#639 opened Dec 20, 2025 by copybara-service bot Loading…
Fix slicing of PRNG keys for jax.pmap usage.
#638 opened Dec 18, 2025 by copybara-service bot Loading…
Move common Shakespeare code to a sharable module to share across examples
#637 opened Dec 18, 2025 by copybara-service bot Loading…
[JAX SC] Update MergeAll method to take an absl::Span<PartitionedCooTensors>. This allows for more flexible usage and avoids unnecessary moves.
#636 opened Dec 16, 2025 by copybara-service bot Loading…
Add table metadata to SparseDenseMatmulConfig.
#615 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…
[JAX SC] Move benchmarks to benchmarks subfolder, tests to separate subfolder.
#602 opened Dec 2, 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! Updated in the last three days: updated:>2025-12-18.