Add embedding shape analysis (#64323)
Summary: Pull Request resolved: https://github.com/pytorch/pytorch/pull/64323
Test Plan: Imported from OSS
Reviewed By: driazati
Differential Revision: D30738145
Pulled By: eellison
fbshipit-source-id: be12408330d671bc65cf645aa2c20fafd954e6a9
