check tensor attrs of tensor wrapper subclasses in prologue

thunder/__init__.py

@@ -592,7 +592,7 @@ def get_computation_and_inputs(*args, **kwargs):
 
             prologue_traces += transform_for_execution(
                 prologue_trc,
-                executors_list=(pythonex,),
+                executors_list=(pythonex, pytorch_executor),
                 use_del_last_used=False,
             )
             prologue_trc = prologue_traces[-1]

thunder/clang/__init__.py

@@ -20,15 +20,16 @@
 from thunder.core import utils
 import thunder.core.prims as prims
 from thunder.core.proxies import (
+    AnyProxy,
     IntegerProxy,
-    NumberProxy,
     NumberLike,
+    NumberProxy,
+    Proxy,
+    SubclassTensorProxy,
     TensorProxy,
-    pyval,
-    pytype,
     proxy,
-    AnyProxy,
-    Proxy,
+    pytype,
+    pyval,
 )
 import thunder.core.devices as devices
 
@@ -67,7 +68,7 @@ def __call__(self, fn: Callable) -> Callable:
 
 # Checks a tensor's shape and metadata (for use with cache check)
 @clangop()
-def check_tensor_shape_and_metadata(t: TensorProxy, /) -> None:
+def check_tensor_shape_and_metadata(t: TensorProxy | SubclassTensorProxy, /) -> None:
     return prims.check_tensor_shape_and_metadata(
         t,
         # replace Proxy entries with `-1`s as wild card, as we any value is

thunder/core/jit_ext.py

@@ -1719,9 +1719,12 @@ def is_variableified_tensorproxy(v: Variable | Proxy) -> Proxy:
 
     with tracectx(prologue_trace):
         for prim, *args in ctx._constraints:
+            subclass_tensor: SubclassTensorProxy | None = None
             for a in args:
                 if isinstance(a, Proxy):
                     unpack(a)
+                if isinstance(a, SubclassTensorProxy):
+                    subclass_tensor = a
             # unpacking Proxy in TensorProxy.shape which is used in `check_tensor_shape_and_metadata`
             if prim == clang.check_tensor_shape_and_metadata:
                 for s in a.shape:
@@ -1730,6 +1733,13 @@ def is_variableified_tensorproxy(v: Variable | Proxy) -> Proxy:
 
             prim(*args)
 
+            if isinstance(subclass_tensor, SubclassTensorProxy):
+                for t in prims.flatten_tensor_subclass(subclass_tensor):
+                    for s in t.shape:
+                        if isinstance(s, Proxy):
+                            unpack(s)
+                    prim(t)
+
         cache_info = thunder._get_cache_info()
         # assert len of cache info to ensure that we're not missing anything?
         if cache_info:

thunder/torch/__init__.py

@@ -3156,18 +3156,15 @@ def amin(a, /, dim=None, keepdim: bool = False):
 
 # NOTE: Using name `torch_max` to avoid conflict with Python's `max`
 @overload
-def torch_max(a: TensorLike, /) -> TensorLike:
-    ...
+def torch_max(a: TensorLike, /) -> TensorLike: ...
 
 
 @overload
-def torch_max(a: TensorLike, /, dim: NumberLike, keepdim: bool = False) -> tuple[TensorLike, TensorLike]:
-    ...
+def torch_max(a: TensorLike, /, dim: NumberLike, keepdim: bool = False) -> tuple[TensorLike, TensorLike]: ...
 
 
 @overload
-def torch_max(a: TensorLike, b: TensorLike, /) -> TensorLike:
-    ...
+def torch_max(a: TensorLike, b: TensorLike, /) -> TensorLike: ...
 
 
 @torchsymbol(torch.max, is_method=True, method_name="max", id="torch.max")