Teach is_signature_compatible() to dig into similar annotations (#2693)

Summary:
Pull Request resolved: #2693

D68450007 updated some annotations in pytorch. This function wasn't correctly evaluating `typing.Dict[X, Y]` and `dict[X, Y]` as the equivalent.

Reviewed By: izaitsevfb

Differential Revision: D68475380

fbshipit-source-id: 3b71ab41f95e6c20986ebe6fbf6f9cbe3b3d58f9

Author: aorenste

torchrec/schema/utils.py

@@ -8,6 +8,32 @@
 # pyre-strict
 
 import inspect
+import typing
+from typing import Any
+
+
+def _is_annot_compatible(prev: object, curr: object) -> bool:
+    if prev == curr:
+        return True
+
+    if not (prev_origin := typing.get_origin(prev)):
+        return False
+    if not (curr_origin := typing.get_origin(curr)):
+        return False
+
+    if prev_origin != curr_origin:
+        return False
+
+    prev_args = typing.get_args(prev)
+    curr_args = typing.get_args(curr)
+    if len(prev_args) != len(curr_args):
+        return False
+
+    for prev_arg, curr_arg in zip(prev_args, curr_args):
+        if not _is_annot_compatible(prev_arg, curr_arg):
+            return False
+
+    return True
 
 
 def is_signature_compatible(
@@ -84,6 +110,8 @@ def is_signature_compatible(
             return False
 
     # TODO: Account for Union Types?
-    if current_signature.return_annotation != previous_signature.return_annotation:
+    if not _is_annot_compatible(
+        previous_signature.return_annotation, current_signature.return_annotation
+    ):
         return False
     return True