Add unstack_and_unshard for SparseCore#147

copybara-service[bot] merged 1 commit intomainjax-ml/jax-tpu-embedding:mainfrom

test_708384192jax-ml/jax-tpu-embedding:test_708384192Copy head branch name to clipboard