Rate this Page
★ ★ ★ ★ ★

is_leaf_nontensor#

tensordict.is_leaf_nontensor(cls: Type) → bool#

Returns True if a type is a tensor or non-tensor data.

Pass it as is_leaf to keys() and the methods that walk the nested entries, to list the non-tensor entries (NonTensorData, NonTensorStack) with the tensors. A UnbatchedTensor is a leaf too. Other tensordicts and tensorclasses are not leaves, and nor is a type that is neither a tensor nor a tensor collection. See Non-tensor data and unbatched entries.

Parameters:

cls (type) – the type of an entry, as returned by entry_class().

Examples

>>> import torch
>>> from tensordict import TensorDict, NonTensorData, is_leaf_nontensor
>>> is_leaf_nontensor(torch.Tensor), is_leaf_nontensor(NonTensorData), is_leaf_nontensor(TensorDict)
(True, True, False)
>>> td = TensorDict(a={}, b="a string!", c=torch.zeros(()))
>>> list(td.keys(leaves_only=True, is_leaf=is_leaf_nontensor))
['b', 'c']

See also

default_is_leaf().