reference/generated/tensordict.TensorDictBase.attrs

Run in Google Colab

Colab

Download Notebook

Notebook

View on GitHub

GitHub

TensorDictBase.attrs

TensorDictBase.attrs(*, fields: Sequence[str] = ('device', 'dtype', 'shape'), num_threads: int | None = None) → Self

Return a deviceless tensordict whose leaves are TensorAttrs.

Each tensor leaf of self is replaced by a TensorAttrs describing the requested tensor attributes. The result is intended to be passed to to() to drive per-leaf device/dtype casting when the source tensordict is heterogeneous.

Keyword Arguments:

• fields (sequence of str, optional) – which attributes to record on each TensorAttrs. Accepts any subset of ("device", "dtype", "shape"). Attributes not listed remain None. Defaults to ("device", "dtype", "shape").

• num_threads (int or None, optional) – number of threads to use when iterating leaves. Defaults to None (single-threaded). Construction of TensorAttrs is Python-bound, so threading typically yields little; exposed for symmetry with to().

Examples

>>> import torch
>>> from tensordict import TensorDict
>>> td = TensorDict(
...     {"a": torch.zeros(3, device="cpu"),
...      "b": torch.zeros(3, dtype=torch.int32)},
...     batch_size=[3],
... )
>>> attrs = td.attrs(fields=("device", "dtype"))
>>> target = TensorDict(a=torch.zeros(3, device="cpu"), b=torch.zeros(3), batch_size=[3])
>>> out = target.to(attrs)   # casts `b` to int32 per-leaf
>>> out["b"].dtype
torch.int32

Every tensordict class and every tensorclass has this method. For the other methods, see Methods by task.