[JIT] add other module apis (#34106)
Summary: Pull Request resolved: https://github.com/pytorch/pytorch/pull/34106
Test Plan: Imported from OSS
Differential Revision: D20283996
Pulled By: eellison
fbshipit-source-id: 88e7bc4547e96717d6c8efe0b25ede0d198d9e68
