-
Notifications
You must be signed in to change notification settings - Fork 3
Pull requests: jax-ml/jax-tpu-embedding
Author
Label
Projects
Milestones
Reviews
Assignee
Sort
Pull requests list
Refactor: Parallelize input preprocessing for minibatching.
#467
opened Sep 13, 2025 by
copybara-service
bot
Loading…
Update the local sparse dense matmul in jax to use the new custom call target, this will ensure that all the gather + scatter run on SparseCore.
#465
opened Sep 11, 2025 by
copybara-service
bot
Loading…
[JAX SC] Add minibatching tests for forward and backward passes.
#462
opened Sep 10, 2025 by
copybara-service
bot
Loading…
Access
EmbeddingVariables table instead of relying on indexing.
#302
opened Jun 17, 2025 by
copybara-service
bot
Loading…
Added fallback logic for computing COO buffer size
#294
opened Jun 12, 2025 by
copybara-service
bot
Loading…
Fix incorrect number of return values in docstring of embedding.preprocess_sparsecore_inputs.
#198
opened Feb 27, 2025 by
copybara-service
bot
Loading…
[CI] Create tpu_presubmit.yml - testing
CI Connection Halt - On Retry
Flag every job that has its workflow file setup for halting to halt if the workflow is retried
#180
opened Feb 7, 2025 by
MichaelHudgins
•
Draft
Split several shared functions into utility modules.
#159
opened Jan 21, 2025 by
copybara-service
bot
Loading…
Input preprocessing library to support vocab-dimension mini-batching.
#155
opened Jan 18, 2025 by
copybara-service
bot
Loading…
Input preprocessing library to support vocab-dimension mini-batching.
#154
opened Jan 18, 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.