Repository: pytorch/tensordict

Issue:
[BUG] Global "indexing" of tensorclasses types leads to type conflicts
debugging test failures from #293 we found out that the solution based on [CLASSES_DICT](https://github.com/pytorch-labs/tensordict/blob/8ab67e42dadd84770452987399efe4f481d9337e/tensordict/prototype/tensorclass.py#L40)
leads to weird behaviour when classes with the same name are being defined (even if locally)

getting the type directly from the instatiated object instead of using the annotation might solve the problem



Discussion:
Thanks for reporting i'll look into this
Can you provide a minimally reproducible example for me to work on?
I am already working on this

Resolving patch:
diff --git a/tensordict/prototype/tensorclass.py b/tensordict/prototype/tensorclass.py
--- a/tensordict/prototype/tensorclass.py
+++ b/tensordict/prototype/tensorclass.py
@@ -11,12 +11,11 @@
 import numbers
 import re
 import sys
-import typing
 import warnings
 from copy import copy
 from dataclasses import dataclass
 from textwrap import indent
-from typing import Any, Callable, Sequence, TypeVar, Union
+from typing import Any, Callable, Sequence, TypeVar
 
 import tensordict as tensordict_lib
 
@@ -36,9 +35,6 @@
 T = TypeVar("T", bound=TensorDictBase)
 PY37 = sys.version_info < (3, 8)
 
-# We keep a dict of str -> class to call the class based on the string
-CLASSES_DICT: dict[str, type] = {}
-
 # Regex precompiled patterns
 OPTIONAL_PATTERN = re.compile(r"Optional\[(.*?)\]")
 UNION_PATTERN = re.compile(r"Union\[(.*?)\]")
@@ -66,7 +62,7 @@
 def is_tensorclass(obj: type | Any) -> bool:
     """Returns True if obj is either a tensorclass or an instance of a tensorclass."""
     cls = obj if isinstance(obj, type) else type(obj)
-    return dataclasses.is_dataclass(cls) and cls.__name__ in CLASSES_DICT
+    return dataclasses.is_dataclass(cls) and hasattr(cls, "__TENSORCLASS")
 
 
 def tensorclass(cls: T) -> T:
@@ -187,10 +183,10 @@ def __torch_function__(
     cls.to_tensordict = _to_tensordict
     cls.device = property(_device, _device_setter)
     cls.batch_size = property(_batch_size, _batch_size_setter)
+    cls.__TENSORCLASS = True
 
     cls.__doc__ = f"{cls.__name__}{inspect.signature(cls)}"
 
-    CLASSES_DICT[cls.__name__] = cls
     tensordict_lib.tensordict._ACCEPTED_CLASSES += [cls]
     return cls
 
@@ -365,8 +361,6 @@ def wrapper(self, item: str) -> Any:
                 and item in self.__dict__["_tensordict"].keys()
             ):
                 out = self._tensordict[item]
-                expected_type = self.__dataclass_fields__[item].type
-                out = _get_typed_output(out, expected_type)
                 return out
             elif (
                 "_non_tensordict" in self.__dict__
@@ -770,28 +764,6 @@ def __ne__(self, other: object) -> bool:
     return _from_tensordict_with_none(self, tensor)
 
 
-def _get_typed_output(out, expected_type: str | type):
-    # from __future__ import annotations turns types in strings. For those we use CLASSES_DICT.
-    # Otherwise, if the output is some TensorDictBase subclass, we check the type and if it
-    # does not match, we map it. In all other cases, just return what has been gathered.
-    if isinstance(expected_type, str) and expected_type in CLASSES_DICT:
-        if isinstance(out, CLASSES_DICT[expected_type]):
-            return out
-        out = CLASSES_DICT[expected_type]._from_tensordict(out)
-    elif (
-        isinstance(expected_type, type)
-        and not isinstance(out, expected_type)
-        and isinstance(out, TensorDictBase)
-    ):
-        out = expected_type._from_tensordict(out)
-    elif isinstance(out, TensorDictBase):
-        dest_dtype = _check_td_out_type(expected_type)
-        if dest_dtype is not None:
-            out = dest_dtype._from_tensordict(out)
-
-    return out
-
-
 def _single_td_field_as_str(key, item, tensordict):
     """Returns a string as a  key-value pair of tensordict.
 
@@ -843,79 +815,3 @@ def _all_non_td_fields_as_str(src_dict) -> list:
             result.append(f"{key}={repr(val)}")
 
     return result
-
-
-def _check_td_out_type(field_def):
-    """This function determines the type of attributes in the tensorclass.
-
-    in order that results from calls to the underlying tensordict
-    can be cast to the expected type before being returned
-    """
-    if PY37:
-        field_def = str(field_def)
-    if isinstance(field_def, str):
-        optional_match = OPTIONAL_PATTERN.search(field_def)
-        union_match = UNION_PATTERN.search(field_def)
-        if optional_match is not None:
-            args = [optional_match.group(1)]
-        elif union_match is not None:
-            args = union_match.group(1).split(", ")
-        else:
-            args = None
-        if args:
-            args = [arg for arg in args if arg not in ("NoneType",)]
-        # skip all Any or TensorDict or Optional[TensorDict] or Union[TensorDict] or Optional[Any]
-        if (args is None and (field_def == "Any" or "TensorDict" in field_def)) or (
-            args and len(args) == 1 and args[0] == "Any"
-        ):
-            return None
-        if args and len(args) == 1 and "TensorDict" in args[0]:
-            return None
-        elif args:
-            # remove the NoneType from args
-            if len(args) == 1 and args[0] in CLASSES_DICT:
-                return CLASSES_DICT[args[0]]
-            if len(args) == 1 and ("TensorDict" in args[0] or args[0] == "Any"):
-                return None
-            else:
-                raise TypeError(
-                    f"{field_def} has args {args} which can't be deterministically cast."
-                )
-        elif args is None:
-            return None
-        else:
-            raise TypeError(
-                f"{field_def} has args {args} which can't be deterministically cast."
-            )
-    else:
-        if typing.get_origin(field_def) is Union:
-            args = typing.get_args(field_def)
-            # remove the NoneType from args
-            args = [arg for arg in args if arg not in (type(None),)]
-            if len(args) == 1 and (
-                typing.Any is not args[0]
-                and args[0] is not TensorDictBase
-                and TensorDictBase not in args[0].__bases__
-            ):
-                # If there is just one type in Union or Optional, we return that type
-                return args[0]
-            elif len(args) == 1 and (
-                typing.Any is args[0]
-                or args[0] is TensorDictBase
-                or TensorDictBase in args[0].__bases__
-            ):
-                # Any or any TensorDictBase subclass are always ok if alone
-                return None
-            else:
-                raise TypeError(
-                    f"{field_def} has args {args} which can't be deterministically cast."
-                )
-        elif (
-            field_def is typing.Any
-            or field_def is TensorDictBase
-            or TensorDictBase in field_def.__bases__
-        ):
-            # Any or any TensorDictBase subclass are alway ok if alone
-            return None
-        else:
-            raise TypeError(f"{field_def} can't be deterministically cast.")


Test patch: