Add flatten loops transformation (#46365)
Summary:
This diff removes the dependency of flattening on tensors by performing flattening on loops instead.
Pull Request resolved: https://github.com/pytorch/pytorch/pull/46365
Reviewed By: ailzhang
Differential Revision: D24366347
Pulled By: navahgar
fbshipit-source-id: 4ba182f37212b6e4033cae13f8e75bc5144389f4
