typing compile_fx.py (#138033)

Summary:
Type annotations for compile_fx.
- Some of the stuff here is pretty complicated (functions which return functions that take functions) so I bailed on those and used `Any` just to get the rest landed.
- There are also changes to type signatures in other files which I did just to let mypy know more about the types in compile_fx.py.

X-link: pytorch/pytorch#138033
Approved by: https://github.com/Skylion007

Reviewed By: wdvr

Differential Revision: D64714765

Pulled By: aorenste

fbshipit-source-id: 262f5cb9b2171e96ce9f895772bd5778ddb4ebe0

userbenchmark/dynamo/dynamobench/_dynamo/utils.py

@@ -516,7 +516,7 @@ def count_calls(g: fx.Graph) -> int:
     return c
 
 
-def identity(x):
+def identity(x: T) -> T:
     return x