Fix some mypy issues in functorch/codegen/gen_vmap_plumbing.py (#82314)
Necessary to make lint happy in the next PR in the stack
Test Plan
- `mypy functorch/codegen/gen_vmap_plumbing.py`
Pull Request resolved: https://github.com/pytorch/pytorch/pull/82314
Approved by: https://github.com/ezyang
