NestedTensor Softmax (#83435)
Summary: Simple mask compute and softmax
Test Plan: unit test
Differential Revision: D38711915
Pull Request resolved: https://github.com/pytorch/pytorch/pull/83435
Approved by: https://github.com/erichan1, https://github.com/huydhn