[Refactor] Better compile checks

ghstack-source-id: c6a8d45
Pull Request resolved: #1139

Author: Vincent Moens

tensordict/_td.py

@@ -98,9 +98,9 @@
 
     _has_funcdim = False
 try:
-    from torch.compiler import is_dynamo_compiling
+    from torch.compiler import is_compiling
 except ImportError:  # torch 2.0
-    from torch._dynamo import is_compiling as is_dynamo_compiling
+    from torch._dynamo import is_compiling
 
 try:
     from torch.nn.parameter import Buffer
@@ -251,7 +251,7 @@ def __init__(
 
         self._tensordict = _StringOnlyDict()
 
-        # if names and is_dynamo_compiling():
+        # if names and is_compiling():
         #     graph_break()
         has_device = device is not None
         sub_non_blocking = False
@@ -284,7 +284,7 @@ def __init__(
                 )
             self._batch_size = self._parse_batch_size(source, batch_size)
             # TODO: this breaks when stacking tensorclasses with dynamo
-            if not is_dynamo_compiling():
+            if not is_compiling():
                 self.names = names
 
             for key, value in source.items():
@@ -313,7 +313,7 @@ def _new_unsafe(
         nested: bool = True,
         **kwargs: dict[str, Any] | None,
     ) -> TensorDict:
-        if is_dynamo_compiling():
+        if is_compiling():
             return TensorDict(
                 source,
                 batch_size=batch_size,
@@ -473,7 +473,7 @@ def _to_module(
         is_dynamo: bool | None = None,
     ):
         if is_dynamo is None:
-            is_dynamo = is_dynamo_compiling()
+            is_dynamo = is_compiling()
         if is_dynamo:
             _check_inbuild()
 
@@ -2264,7 +2264,7 @@ def _parse_batch_size(
     ) -> torch.Size:
         ERR = "batch size was not specified when creating the TensorDict instance and it could not be retrieved from source."
 
-        if is_dynamo_compiling():
+        if is_compiling():
             if isinstance(batch_size, torch.Size):
                 return batch_size
             elif isinstance(batch_size, tuple):
@@ -2316,7 +2316,7 @@ def names(self):
 
     @names.setter
     def names(self, value):
-        if is_dynamo_compiling():
+        if is_compiling():
             if value is not None:
                 graph_break()
             else:

tensordict/_torch_func.py

@@ -39,9 +39,9 @@
 )
 
 try:
-    from torch.compiler import is_dynamo_compiling
+    from torch.compiler import is_compiling
 except ImportError:  # torch 2.0
-    from torch._dynamo import is_compiling as is_dynamo_compiling
+    from torch._dynamo import is_compiling
 
 TD_HANDLED_FUNCTIONS: dict[Callable, Callable] = {}
 LAZY_TD_HANDLED_FUNCTIONS: dict[Callable, Callable] = {}
@@ -301,7 +301,7 @@ def _cat(
         out = {}
         for key in keys:
             items = [td._get_str(key, NO_DEFAULT) for td in list_of_tensordicts]
-            if not is_dynamo_compiling():
+            if not is_compiling():
                 with _ErrorInteceptor(
                     key, "Attempted to concatenate tensors on different devices at key"
                 ):
@@ -335,7 +335,7 @@ def _cat(
                 _ErrorInteceptor(
                     key, "Attempted to concatenate tensors on different devices at key"
                 )
-                if not is_dynamo_compiling()
+                if not is_compiling()
                 else contextlib.nullcontext()
             ):
                 if isinstance(out, TensorDict):
@@ -592,7 +592,7 @@ def stack_fn(key, values, is_not_init, is_tensor):
                     _ErrorInteceptor(
                         key, "Attempted to stack tensors on different devices at key"
                     )
-                    if not is_dynamo_compiling()
+                    if not is_compiling()
                     else contextlib.nullcontext()
                 ):
                     return _stack(values, dim, maybe_dense_stack=maybe_dense_stack)

tensordict/base.py

@@ -103,9 +103,9 @@
 from torch.utils._pytree import tree_map
 
 try:
-    from torch.compiler import is_dynamo_compiling
+    from torch.compiler import is_compiling
 except ImportError:  # torch 2.0
-    from torch._dynamo import is_compiling as is_dynamo_compiling
+    from torch._dynamo import is_compiling
 
 try:
     from torch import _foreach_copy_
@@ -5247,7 +5247,7 @@ def _view_and_pad(tensor):
                 if v.device != storage.device:
                     v = v.to(storage.device, non_blocking=non_blocking)
                 stride = v.stride()
-                if is_dynamo_compiling():
+                if is_compiling():
                     if not v.is_contiguous():
                         v = v.clone(memory_format=torch.contiguous_format)
                 elif (stride and stride[-1] != 1) or v.storage_offset():
@@ -6963,7 +6963,7 @@ def _values_list(
                 is_leaf=is_leaf,
                 collapse=collapse,
             )
-            if is_dynamo_compiling():
+            if is_compiling():
                 key_to_index = {key: i for i, key in enumerate(keys)}
                 return [vals[key_to_index[key]] for key in sorting_keys]
             else:
@@ -6994,7 +6994,7 @@ def _items_list(
             return list(keys), list(vals)
         if default is None:
             # TODO: check that lists are identical
-            if is_dynamo_compiling():
+            if is_compiling():
                 key_to_index = {key: i for i, key in enumerate(keys)}
                 new_vals = [vals[key_to_index[key]] for key in sorting_keys]
                 if len(new_vals) < len(vals):
@@ -7015,12 +7015,9 @@ def _items_list(
             ]  # intersection does not keep the sorting
         else:
             new_keys = list(set(sorting_keys).union(keys))
-        if is_dynamo_compiling():
-            ...
-        else:
-            source = dict(zip(keys, vals))
-            vals = [source.get(key, default) for key in new_keys]
-            return new_keys, vals
+        source = dict(zip(keys, vals))
+        vals = [source.get(key, default) for key in new_keys]
+        return new_keys, vals
 
     def _grad(self):
         # We can't cache this because zero_grad can be called outside (eg from optimizer) and we want the tensors
@@ -11931,7 +11928,7 @@ def unflatten_keys(self, separator: str = ".", inplace: bool = False) -> T:
                 result.lock_()
             return result
         else:
-            if not is_dynamo_compiling():
+            if not is_compiling():
                 key_list = list(self.keys())
             else:
                 key_list = [k for k in self.keys()]  # noqa
@@ -12196,10 +12193,10 @@ def lock_(self) -> T:
         """
         if self.is_locked:
             return self
-        is_compiling = is_dynamo_compiling()
-        if is_compiling:
+        is_comp = is_compiling()
+        if is_comp:
             _lock_warn()
-        self._propagate_lock(is_compiling=is_compiling)
+        self._propagate_lock(is_compiling=is_comp)
         return self
 
     @erase_cache
@@ -12611,7 +12608,7 @@ def copy_dict(d):
     def _sync_all(self):
         if _has_cuda:
             # TODO: dynamo doesn't like torch.cuda.is_initialized
-            if not is_dynamo_compiling() and torch.cuda.is_initialized():
+            if not is_compiling() and torch.cuda.is_initialized():
                 torch.cuda.synchronize()
         elif _has_mps:
             mps = getattr(torch, "mps", None)
@@ -12799,7 +12796,7 @@ def _register_tensor_class(cls):
 
 
 def _is_tensor_collection(datatype: type) -> bool:
-    is_dynamo = is_dynamo_compiling()
+    is_dynamo = is_compiling()
     out = None
     if not is_dynamo:
         out = _TENSOR_COLLECTION_MEMO.get(datatype)

tensordict/nn/common.py

@@ -27,9 +27,9 @@
 from torch import nn, Tensor
 
 try:
-    from torch.compiler import is_dynamo_compiling
+    from torch.compiler import is_compiling
 except ImportError:  # torch 2.0
-    from torch._dynamo import is_compiling as is_dynamo_compiling
+    from torch._dynamo import is_compiling
 
 try:
     from functorch import FunctionalModule, FunctionalModuleWithBuffers
@@ -1153,7 +1153,7 @@ def __repr__(self) -> str:
         return f"{type(self).__name__}(\n{fields})"
 
     def __getattr__(self, name: str) -> Any:
-        if not is_dynamo_compiling():
+        if not is_compiling():
             __dict__ = self.__dict__
             _parameters = __dict__.get("_parameters")
             if _parameters:
@@ -1230,7 +1230,7 @@ def __init__(self, td_module: TensorDictModuleBase) -> None:
                 self.register_forward_hook(self.td_module._forward_hooks[pre_hook])
 
     def __getattr__(self, name: str) -> Any:
-        if not is_dynamo_compiling():
+        if not is_compiling():
             __dict__ = self.__dict__
             _parameters = __dict__.get("_parameters")
             if _parameters:

tensordict/nn/utils.py

@@ -18,9 +18,9 @@
 from torch.utils._contextlib import _DecoratorContextManager
 
 try:
-    from torch.compiler import is_dynamo_compiling
+    from torch.compiler import is_compiling
 except ImportError:  # torch 2.0
-    from torch._dynamo import is_compiling as is_dynamo_compiling
+    from torch._dynamo import is_compiling
 
 
 _dispatch_tdnn_modules = _ContextManager(
@@ -300,7 +300,7 @@ def wrapper(_self, tensordict, *args: Any, **kwargs: Any) -> Any:
         return super().__call__(wrapper)
 
     def __enter__(self) -> None:
-        if self.mode and is_dynamo_compiling():
+        if self.mode and is_compiling():
             raise RuntimeError("skip_existing is not compatible with TorchDynamo.")
         self.prev = _skip_existing.get_mode()
         if self.mode is not None:
@@ -338,7 +338,7 @@ def __call__(self, func: Callable):
 
         @functools.wraps(func)
         def wrapper(_self, tensordict, *args: Any, **kwargs: Any) -> Any:
-            if skip_existing() and is_dynamo_compiling():
+            if skip_existing() and is_compiling():
                 raise RuntimeError(
                     "skip_existing is not compatible with torch.compile."
                 )
@@ -351,7 +351,7 @@ def wrapper(_self, tensordict, *args: Any, **kwargs: Any) -> Any:
                 and not any(key in out_keys for key in in_keys)
             ):
                 return tensordict
-            if is_dynamo_compiling():
+            if is_compiling():
                 return func(_self, tensordict, *args, **kwargs)
             self.prev = _skip_existing.get_mode()
             try:

tensordict/tensorclass.py

@@ -64,9 +64,9 @@
 from torch.utils._pytree import tree_map
 
 try:
-    from torch.compiler import is_dynamo_compiling
+    from torch.compiler import is_compiling
 except ImportError:  # torch 2.0
-    from torch._dynamo import is_compiling as is_dynamo_compiling
+    from torch._dynamo import is_compiling
 
 
 def _identity(cls):
@@ -890,7 +890,7 @@ def wrapper(
         if lock is None:
             lock = frozen
 
-        if not is_dynamo_compiling():
+        if not is_compiling():
             # zip not supported by dynamo
             for value, key in zip(args, self.__dataclass_fields__):
                 if key in kwargs:
@@ -904,7 +904,7 @@ def wrapper(
 
         if batch_size is None:
             batch_size = torch.Size([])
-        if not is_dynamo_compiling():
+        if not is_compiling():
             for key, field in type(self).__dataclass_fields__.items():
                 if field.default_factory is not dataclasses.MISSING:
                     default = field.default_factory()
@@ -1072,7 +1072,7 @@ def _from_tensordict(cls, tensordict, non_tensordict=None, safe=True):  # noqa:
     # tensordict = tensordict.copy()
     tensor_keys = tensordict.keys()
     # TODO: compile doesn't like set() over an arbitrary object
-    if is_dynamo_compiling():
+    if is_compiling():
         tensor_keys = {k for k in tensor_keys}  # noqa: C416
         exp_keys = {k for k in cls.__expected_keys__}  # noqa: C416
         if non_tensordict is not None:
@@ -1112,7 +1112,7 @@ def _from_tensordict(cls, tensordict, non_tensordict=None, safe=True):  # noqa:
         for key in to_add:
             non_tensordict[key] = None
 
-    if not is_dynamo_compiling():
+    if not is_compiling():
         # bypass initialisation. this means we don't incur any overhead creating an
         # empty tensordict and writing values to it. we can skip this because we already
         # have a tensordict to use as the underlying tensordict
@@ -1313,7 +1313,7 @@ def wrapper(self, key: str, value: Any) -> None:  # noqa: D417
             value (any): the value to set for the attribute
 
         """
-        if not is_dynamo_compiling():
+        if not is_compiling():
             __dict__ = self.__dict__
             if (
                 "_tensordict" not in __dict__
@@ -1348,7 +1348,7 @@ def deliver_result(self, result, kwargs):
         if result is None:
             return
         if isinstance(result, TensorDictBase) and kwargs.get("out") is not result:
-            if not is_dynamo_compiling():
+            if not is_compiling():
                 non_tensordict = super(type(self), self).__getattribute__(
                     "_non_tensordict"
                 )
@@ -1362,7 +1362,7 @@ def deliver_result(self, result, kwargs):
         return result
 
     def wrapped_func(self, *args, **kwargs):
-        if not is_dynamo_compiling():
+        if not is_compiling():
             td = super(type(self), self).__getattribute__("_tensordict")
         else:
             td = self._tensordict
@@ -1409,7 +1409,7 @@ def wrapped_func(*args, **kwargs):
             return type(self)._from_tensordict(res, dict(self._non_tensordict))
         return res
 
-    if not is_dynamo_compiling():
+    if not is_compiling():
         wrapped_func = functools.wraps(func)(wrapped_func)
 
     return wrapped_func