[BUG] TensorClass with non_tensor lead to torch.compile failing (RecursionError or compiling forever)

[mctigger]

Describe the bug

When a @tensorclass object is created inside a compile function and the tensorclass has fields that are not a torch.Tensor (enforced through either being of a type that cannot be cast to tensor or with nocast=True), torch.compile will run forever or throw an error. The code below explains the issue the best.

To Reproduce

Note: The cases may have different failure behavior as stated in the code. With tensordict-nightly==02025.2.3 I produced the one below. For tensordict 0.8.3 there is different behavior, but still wrong behavior (e.g. compiling forever).

import torch
from tensordict import TensorClass
from torch import Tensor


# Same issue with the decorator @tensorclass version
class TensorClassWithNonTensorData(TensorClass["nocast"]):
    tensor: Tensor
    non_tensor_data: int


data = torch.ones(4, 5, device="cuda")


def fn_no_device_no_batch_size():
    a = TensorClassWithNonTensorData(tensor=data, non_tensor_data=1)
    return a.tensor


def fn_no_device():
    a = TensorClassWithNonTensorData(tensor=data, non_tensor_data=1, batch_size=[4])
    return a.tensor


def fn_with_device():
    a = TensorClassWithNonTensorData(
        tensor=data, non_tensor_data=1, batch_size=[4], device="cuda"
    )
    return a.tensor


def fn_with_device_without_batch_size():
    a = TensorClassWithNonTensorData(tensor=data, non_tensor_data=1, device="cuda")
    return a.tensor


mode = "with_device_without_batch_size"

match mode:
    case "no_device_no_batch_size":
        # This will print
        print("no_device_no_batch_size", torch.compile(fn_no_device_no_batch_size)())
    case "no_device":
        # This throws an error. Interestingly, if the previous case is executed before,
        # then this doesn't throw an error!
        print("no_device", torch.compile(fn_no_device)())
    case "with_device":
        # This throws an error! Interestinfly, if the first case is executed before,
        # this does not throw an error, but compiles forever.
        print("with_device", torch.compile(fn_with_device)())
    case "with_device_without_batch_size":
        # This compiles forever, no matter what"
        print(
            "with_device_without_batch_size",
            torch.compile(fn_with_device_without_batch_size)(),
        )

This is the error if an error is thrown (for some cases it compiles forever, see above).

torch._dynamo.exc.InternalTorchDynamoError: RecursionError: maximum recursion depth exceeded while calling a Python object

from user code:
   File "/opt/conda/lib/python3.10/site-packages/tensordict/_td.py", line 2453, in _set_str
    value = self._validate_value(
  File "/opt/conda/lib/python3.10/site-packages/tensordict/base.py", line 11108, in _validate_value
    value.batch_size = self.batch_size
  File "/opt/conda/lib/python3.10/site-packages/tensordict/tensorclass.py", line 1530, in wrapper
    return setattr_(self, key, value)
  File "/opt/conda/lib/python3.10/site-packages/tensordict/tensorclass.py", line 1530, in wrapper
    return setattr_(self, key, value)
  File "/opt/conda/lib/python3.10/site-packages/tensordict/tensorclass.py", line 1530, in wrapper
    return setattr_(self, key, value)
  [Previous line repeated 64 more times]

Expected behavior

It should compile.

System info

Describe the characteristic of your environment:

• Python 3.10.16
• tensordict-nightly==02025.2.3
• torch==2.6.0

Checklist

• I have checked that there is no similar issue in the repo (required)
• I have read the documentation (required)
• I have provided a minimal working example to reproduce the bug (required)

[vmoens]

On it!

[mctigger]

Hi @vmoens thank you very much for reacting so quickly and fixing the compile bug!

Please correct me if I am wrong, but just setting the fields to NonTensorData ignores MetaData fields.
In my case, I use tensorclass to also hold some torch.Size data that is essentially meta data that does not have a batch size and is not transformable (through stack or view or anything).

Previously, this was kept as it is, but now it seems to also expand with the TensorClass, breaking my code.
Should we maybe check whether it is MetaData?
And also I would nocast=True expected to put everything that is not a tensor into MetaData, i.e. not only not be a tensor, but also to not behave like a tensor. Maybe we can either change this or add a sentence to the tensorclass doc explaining this.

[vmoens]

Previously, this was kept as it is, but now it seems to also expand with the TensorClass, breaking my code.
Should we maybe check whether it is MetaData?
Yeah so you're right that the behaviour has changed and it's not ideal.

Why the bc-breaking change

Mainly because it resolves a bug. Previously, a single class was handling both metadata and batched non-tensor data. With LLM, we often have things like

td = TensorDict(a=['a string'], batch_size=(1,))
td.squeeze().unsqueeze(0)

With NonTensorData, the squeeze gives you a tensordict with NonTensorData with no batch-size and unsqueeze just adds a batch size but you don't get back a list of strings like you had before. Hence, squeeze/unsqueeze is not bijective and that's very hard to solve!

That looked like the most "common" behaviour (ie, the behaviour I see the most often lol).

So I created another substitute for this called MetaData.

I can revert these changes and make the new behaviour optional through a context manager / env variable or whatever it is that we want to use to control that globally. That looked like one more thing of that kind and we already have a bunch so I was considering just making the BC-breaking change and documenting it. I will rely on what the community tells me though!

So the way to go for you would be to wrap you non-tensor data manually in a MetaData (not ideal, I agree, but we must find a way to accommodate everyone!) Do you have another suggestion to take care of this? I'm all ears!

And also I would nocast=True expected to put everything that is not a tensor into MetaData, i.e. not only not be a tensor, but also to not behave like a tensor.

You mean "everything that is not a tensor would be a MetaData and not a NonTensorData? The indended usage is also to carry lists of strings that are batched with LLMs. If we don't use nocast`, then TensorDict will put your strings in a numpy array I think, which I don't really like.

Anyway, thanks for the feedback. Let's try to smooth this out before the 0.9 release! I defo don't want people to struggle to move from 0.8 to 0.9!

[mctigger]

Thanks for again, answering so quickly and taking the time to explain.

I am agreeing with you on all points and I think proceeding as you say is reasonable.

However, I am wondering how to get the behavior I need: Just assigning an attribute to some value that stays constant all the time.

Something like:

class FixedAttributeTensorClass(TensorClass):
    my_string: Str

a = FixedAttributeTensorClass("Hello")
b = FixedAttributeTensorClass("Hello")

c = torch.stack([a,b])

# Should print "Hello"
print(c.my_string)

And for different meta data:

a = FixedAttributeTensorClass("This")
b = FixedAttributeTensorClass("is not the same")

# Should throw exception because meta data does not match
c = torch.stack([a,b])

Maybe I am overlooking something, but this does not seem to be possible at the moment?

[vmoens]

Previously we were checking equalities but it was tedious and error prone (what if your class does not have an __eq__?) - however, we could find a way to make it work (see below).

I think I have a solution: Using typing (as with pydantic)

Here is an example. Essentially, if you tell me x: str then I consider it's metadata, if you tell me x: NestedList[int] (name tbd) I consider it's NonTensorData + stack

Prototype

from typing import Any, Generic, List, Sequence, TypeVar, Union, get_args, get_origin

import torch
from pydantic import BaseModel, model_validator
from tensordict import (
    MetaData,
    NonTensorData,
    NonTensorStack,
    TensorDict,
    is_tensor_collection,
    set_list_to_stack,
)
from torch import Size, Tensor

set_list_to_stack(True).set()

T = TypeVar("T")


class NestedList(Generic[T]):
    """A type that accepts either a single value or nested lists of values of type T"""

    def __init__(self, value: Union[T, List["NestedList[T]"]]):
        self.value = value

    @classmethod
    def __get_pydantic_core_schema__(cls, source_type, _handler):
        T = get_args(source_type)[0]  # Get the type parameter

        def validate_nested(v):
            if isinstance(v, (list, tuple)):
                return [validate_nested(x) for x in v]
            return T(v) if not isinstance(v, T) else v

        return {"type": "any", "mode": "before", "function": validate_nested}


class TensorClass(BaseModel):
    model_config = {
        "arbitrary_types_allowed": True,
        "extra": "allow",  # Allow extra fields like batch_size
    }

    def __init__(self, **data):
        batch_size = data.pop("batch_size", None)
        device = data.pop("device", None)
        super().__init__(**data)
        self.__dict__["_tensordict"] = TensorDict(
            {}, batch_size=batch_size, device=device
        )
        # Initialize tensordict with current values
        for field in self.__class__.model_fields:
            if hasattr(self, field):
                item = getattr(self, field)
                if isinstance(item, torch.Tensor) or is_tensor_collection(item):
                    self._tensordict[field] = item
                else:
                    # Get the field's type annotation
                    field_type = self.__class__.model_fields[field].annotation

                    def has_nested_list(type_):
                        """Check if a type contains NestedList"""
                        if get_origin(type_) is NestedList:
                            return True
                        if get_origin(type_) in (Union, None):
                            # For Union types or simple types, check each argument
                            args = get_args(type_)
                            return any(has_nested_list(arg) for arg in args)
                        return False

                    def get_primary_type(type_):
                        """Get the primary type (non-NestedList) from a Union or simple type"""
                        if get_origin(type_) in (Union, None):
                            # Look through Union args for non-NestedList type
                            for arg in get_args(type_):
                                if not has_nested_list(arg):
                                    return arg
                        return None

                    # Check if it's a NestedList type or contains NestedList
                    if has_nested_list(field_type):
                        primary_type = get_primary_type(field_type)
                        # If it matches the primary type (e.g. str for b), treat as scalar
                        if primary_type and isinstance(item, primary_type):
                            self._tensordict[field] = NonTensorData(item)
                        # Otherwise if it's a sequence (but not str), treat as list
                        elif isinstance(item, Sequence) and not isinstance(
                            item, (str, bytes)
                        ):
                            stack = NonTensorStack.from_list(item)
                            self._tensordict[field] = stack
                        else:
                            self._tensordict[field] = NonTensorData(item)
                    else:
                        self._tensordict[field] = MetaData(item)
                delattr(self, field)

    @property
    def device(self) -> torch.device:
        return self._tensordict.device

    @device.setter
    def device(self, device: torch.device):
        self._tensordict.device = device

    @property
    def batch_size(self) -> Size:
        """Get the batch size of the underlying TensorDict"""
        return self._tensordict.batch_size

    @batch_size.setter
    def batch_size(self, size):
        """Set the batch size of the underlying TensorDict"""
        td = self.__dict__["_tensordict"]
        td.batch_size = size

    @model_validator(mode="after")
    def sync_to_tensordict(self, data):
        # Ensure _tensordict exists with proper batch size
        if not hasattr(self, "_tensordict"):
            self.__dict__["_tensordict"] = TensorDict({}, batch_size=[])

        # Sync all fields to tensordict
        for field in self.__class__.model_fields:
            if hasattr(self, field):
                self._tensordict[field] = getattr(self, field)
        return self

    def __getattr__(self, name: str) -> Any:
        if name == "_tensordict":
            try:
                return self.__dict__[name]
            except KeyError:
                raise AttributeError(
                    f"{self.__class__.__name__} has no attribute {name}"
                )
        try:
            return self._tensordict[name]
        except (KeyError, AttributeError):
            raise AttributeError(f"{self.__class__.__name__} has no attribute {name}")

    def __repr__(self):
        fields_repr = ", ".join(
            f"{field}={getattr(self, field)}" for field in self.__class__.model_fields
        )
        extra_fields = {
            "batch_size": self.batch_size,
            "device": self._tensordict.device,
        }
        extra_repr = ", ".join(f"{k}={v}" for k, v in extra_fields.items())
        return f"{self.__class__.__name__}({fields_repr}, {extra_repr})"


class MyClass(TensorClass):
    a: int | NestedList[int]  # NonTensorStack
    b: str | NestedList[str]  # Now accepts str or list[str] or list[list[str]] etc
    c: Tensor
    d: str  # MetaData


# Default (empty) batch size
model = MyClass(a=1, b="hello", c=torch.tensor([1.0, 2.0, 3.0]), d="a string")
print(f"{model=}")
print(f"Model attributes: a={model.a}, b={model.b}")
print(f"Model tensor c={model.c}")
print(f"TensorDict contents: {model._tensordict}")
print(f"Default batch_size: {model.batch_size}")

# Integer batch size
model = MyClass(
    a=[1, 2],
    b=["hello", "world"],
    c=torch.tensor([[1.0], [2.0]]),  # 2x1 tensor
    d="a string",
    batch_size=2,
)
print(f"{model=}")
print(f"{model._tensordict=}")
print(f"Integer batch_size: {model.batch_size}")

I think for your example, we should have a way to put in the model_config that we guarantee that the MetaData can be compared (or provide a function to do so?).
PyTorch now has a fast implementation of pytree that we could leverage for that.

[vmoens]

Here is a PR: #1342

[mctigger]

Looks good, thank you very much!