[vmap] Add batching rules for comparisons ops (#50364)
Summary:
Related to https://github.com/pytorch/pytorch/issues/49562
This PR adds batching rules for the below comparison ops.
- torch.eq
- torch.gt
- torch.ge
- torch.le
- torch.lt
- torch.ne
Pull Request resolved: https://github.com/pytorch/pytorch/pull/50364
Reviewed By: anjali411
Differential Revision: D25885359
Pulled By: zou3519
fbshipit-source-id: 58874f24f8d525d8fac9062186b1c9970618ff55
