optax icon indicating copy to clipboard operation
optax copied to clipboard

Published 20 hours ago • verified publisher icon google-deepmind

Support for CSR format sparse matrix in optimizer?

Open MoFHeka opened this issue 7 months ago • 1 comments

matrix format:https://jax.readthedocs.io/en/latest/_autosummary/jax.experimental.sparse.BCSR.html#jax.experimental.sparse.BCSR

some feature in deepspeed: https://github.com/microsoft/DeepSpeed/blob/master/deepspeed/ops/csrc/adam/fused_adam_frontend.cpp

MoFHeka avatar Jun 28 '24 09:06 MoFHeka