-
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
Add the Adam optimizer from [Kingma et al., 2014](http://arxiv.org/abs/1412.6980).
#264
opened May 28, 2025 by
copybara-service
bot
Loading…
Do not count the sentinel node for checking the buffer overflow as it is not actually filled in the COO buffer.
#263
opened May 27, 2025 by
copybara-service
bot
Loading…
add a use shardy true variant for
trainer_main_sparsecore_test
#262
opened May 27, 2025 by
copybara-service
bot
Loading…
Introduce FdoStats struct for tracking FDO info.
#253
opened May 12, 2025 by
copybara-service
bot
Loading…
Refactor to use a config object to manage arguments to sparse dense matmul.
#235
opened Apr 25, 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 p on any issue or pull request to go back to the pull request listing page.