Match AITModule on ait inputs and fx inputs (#262)
Summary:
X-link: https://github.com/facebookincubator/AITemplate/pull/262
Currently AITModule doesn't
check if python's input and AIT input matches, but directly give AIT the corresponding input.
This could cause us some trouble when AIT optimize the graph and eliminate some inputs.
e.g. https://www.internalfb.com/phabricator/paste/view/P619446970?lines=37
Here, add_14 can be deduced at a very early stage and therefore the source of add_14: cat_37 actually isn't used in AIT graph anywhere!
Therefore, we would got complaint: User passed 6 inputs, but the model expects 5. This is actually because cat_37 is not longer needed in AIT's graph!
This diff put checker in AITinterpreter to insure eeach AIT input can find corresponding python's input.
It also change AITModule by checking both list of python/fx input and ait input, and find corresponding fx input for ait input.
Reviewed By: frank-wei, chenyang78
Differential Revision: D43256832
fbshipit-source-id: 369bfddfc0bfba20e62d794b3ed4fdc2797ce9ec
