Generalize expand logic (#63615)
Summary:
Pull Request resolved: https://github.com/pytorch/pytorch/pull/63615
how to review: pretty much just check that the inputs generated are a good representation of the op semantics, that should be sufficient for correctness, and then you can also double check the op size semantics by going to https://codebrowser.bddppq.com/pytorch/pytorch/ typing in native::{op_name} and looking at the op implementation as a bonus if you want
Test Plan: Imported from OSS
Reviewed By: driazati
Differential Revision: D30738148
Pulled By: eellison
fbshipit-source-id: 4ef74a9c9b39c0beb73949e63aa844c46ab637eb