reference/utils

Run in Google Colab

Colab

Download Notebook

Notebook

View on GitHub

GitHub

Utilities

This page lists the names of tensordict.utils that the tensordict module does not export: helpers for tensors and tensordicts, the types that some methods return, and type aliases. Import them from tensordict.utils. The names that tensordict also exports, such as unravel_key() or assert_close(), are on tensordict: containers and functions.

Tensor and tensordict helpers

expand_right and expand_as_right add dimensions on the right of a tensor, where broadcasting adds them on the left. isin and remove_duplicates compare the values of one entry along a batch dimension.

expand_right(tensor, shape)	Expand a tensor on the right to match a desired shape.
expand_as_right(tensor, dest)	Expand a tensor on the right to match another tensor shape.
isin(input, reference, key[, dim])	Tests if each element of key in input dim is also present in the reference.
remove_duplicates(input, key[, dim, ...])	Removes indices duplicated in key along the specified dimension.

Returned types

The methods that have a return_early argument, such as memmap(), return a TensorDictFuture if you pass return_early=True. tolist() returns a LinkedList if you pass as_linked_list=True.

TensorDictFuture(futures, resulting_td)	A custom future class for TensorDict multithreaded operations.
LinkedList(*args, td, **kwargs)	A list of the values of a non-tensor entry that writes item assignments back to the tensordict.

Type aliases

DeviceType is the type of the device arguments, and IndexType is the type of the indices that a tensordict accepts.

DeviceType	alias of Union[device, str, int]
IndexType	alias of Union[None, int, slice, str, Tensor, List[Any], Tuple[Any, ...]]

Debugging and profiling

print_directory_tree prints the files of a directory, for example the directory of a saved tensordict. timeit and logger are the timer and the logger that tensordict uses in its own code. Buffer is torch.nn.Buffer, kept here so that old imports of tensordict.utils.Buffer still work.

print_directory_tree(path[, indent, ...])	Prints the directory tree starting from the specified path.
timeit(name)	Measures the time spent in a function or a block of code, by name.
Buffer([data, persistent])	A kind of Tensor that should not be considered a model parameter.

tensordict.utils.logger

The logging.Logger named "tensordict". To show fewer messages from tensordict, raise the level of this logger, for example logging.getLogger("tensordict").setLevel(logging.WARNING).

See also

• tensordict: containers and functions: the functions of the tensordict module.

• Save and load tensordicts: memmap(return_early=True) and the files of a saved tensordict.