flax
Add vmap reset() regression test for nnx.metrics (#5483)
#5491
Open

Add vmap reset() regression test for nnx.metrics (#5483) #5491

chenkuanliao wants to merge 1 commit into google:main from chenkuanliao:fix-issue-5483
chenkuanliao
chenkuanliao chenkuanliao force pushed from db1487c0 to a71a8f3c 17 hours ago
chenkuanliao Add vmap reset() regression test for nnx.metrics (#5483)
ca14973a
chenkuanliao chenkuanliao force pushed from a71a8f3c to ca14973a 17 hours ago
chenkuanliao chenkuanliao changed the title Make metric reset() shape-preserving under vmap (#5483) Add vmap reset() regression test for nnx.metrics (#5483) 17 hours ago

Login to write a write a comment.

Login via GitHub

Reviewers
No reviews
Assignees
No one assigned
Labels
Milestone