Revert "Cleanup module imports"

This reverts commit 9e1b27571f57f395bb1d33dce276e141ec501cd2.