flax
[nnx] vectorize vmap split counts
#3989
Merged

[nnx] vectorize vmap split counts #3989

copybara-service merged 1 commit into main from nnx-fix-vmap
cgarciae
cgarciae [nnx] vectorize vmap split counts
e4f417ac
cgarciae cgarciae added pull ready
chiamp
chiamp approved these changes on 2024-06-12
copybara-service copybara-service merged 551eaf02 into main 2 years ago
copybara-service copybara-service deleted the nnx-fix-vmap branch 2 years ago

Login to write a write a comment.

Login via GitHub

Reviewers
Assignees
No one assigned
Labels
Milestone