Traceable wrapper subclass support for deferred runtime asserts

torch/fx/experimental/symbolic_shapes.py

@@ -485,6 +485,18 @@ def get(self, o: Any) -> Any:
         return getattr(o, self.name)()
 
 
+@dataclass(frozen=True)
+class InnerTensorKey:
+    inner_name: str
+
+    def __str__(self) -> str:
+        return f".{self.inner_name}"
+
+    def get(self, o: Any) -> Any:
+        """Call the method on object"""
+        return getattr(o, self.inner_name)
+
+
 @dataclass(frozen=True)
 class DivideByKey:
     divisor: int
@@ -534,6 +546,14 @@ def free_unbacked_symbols_with_path(
                             real=real[i] if real is not None else None
                         )
                     )
+            elif is_traceable_wrapper_subclass(a):
+                # TODO: Determine if this is correct
+                attrs, _ = a.__tensor_flatten__()
+                for attr in attrs:
+                    sub = getattr(a, attr)
+                    r.update(
+                        free_unbacked_symbols_with_path(sub, path + (InnerTensorKey(attr),))
+                    )
             elif isinstance(a, torch.Tensor):
                 r.update(
                     free_unbacked_symbols_with_path(

torch/fx/passes/runtime_assert.py

@@ -56,6 +56,7 @@ def insert_deferred_runtime_asserts(
         cast_symbool_to_symint_guardless,
         ConvertIntKey,
         DivideByKey,
+        InnerTensorKey,
         free_symbols,
     )
     from torch.utils._sympy.interp import sympy_interp
@@ -225,6 +226,13 @@ def go(node, keypath):
                                 ),
                                 keypath[1:],
                             )
+                        elif isinstance(keypath[0], InnerTensorKey):
+                            return go(
+                                graph.call_function(
+                                    getattr, (node, keypath[0].inner_name)
+                                ),
+                                keypath[1:],
+                            )
                         else:
                             raise AssertionError(f"unrecognized keypath {keypath}")