-
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
[JAX SC] Add Ragged Tensor Input Batch Format
#325
opened Jun 26, 2025 by
copybara-service
bot
Loading…
[JAX SC] Refactor ragged tensor input batch to reuse common methods and logic.
#324
opened Jun 26, 2025 by
copybara-service
bot
Loading…
[JAX SC] Use Stream API for
SparseCooInputBatch
#323
opened Jun 26, 2025 by
copybara-service
bot
Loading…
[JAX SC] Clean up
tf.sparse.tensor preprocessing lib.
#309
opened Jun 18, 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…
Expose run‑time FDO limits through
FdoParams and InputParams in the API.
#293
opened Jun 12, 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…
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!
Mix and match filters to narrow down what you’re looking for.