flax
Add simple argument-only lifted nn.grad function.
#3463

Merged

Add simple argument-only lifted nn.grad function. #3463

copybara-service merged 1 commit into google:main from levskaya:vjp_fix

levskaya requested a review from

cgarciae 2 years ago

levskaya added pull ready

levskaya force pushed 2 years ago

cgarciae commented on 2023-11-06

levskaya force pushed 2 years ago

cgarciae approved these changes on 2023-11-06

levskaya force pushed 2 years ago

Add simple argument-only lifted nn.grad and nn.value_and_grad functions.

5ff36ba4

levskaya force pushed to 5ff36ba4 2 years ago

copybara-service merged e2c3dfd3 into main 2 years ago

Reviewers

cgarciae

Assignees

No one assigned

Labels

pull ready

Milestone

No milestone