[functorch] follow-up vmapjvpvjp (#84992)
Ref 1: https://github.com/pytorch/pytorch/pull/83375#discussion_r970046113
Ref 2: https://github.com/pytorch/pytorch/pull/83375#discussion_r970047848
Pull Request resolved: https://github.com/pytorch/pytorch/pull/84992
Approved by: https://github.com/zou3519
