tensordict: containers and functions#
This page lists the tensordict containers, and the methods that they all
share, grouped by task. Each group also lists the functions of the
tensordict module for its task. A TensorDict is a
dict[str, torch.Tensor] that behaves like one tensor. Its entries share
their leading dimensions, the batch size, and an operation such as
indexing, reshaping, or to() applies to every
entry. For an introduction, see What is TensorDict?.
Containers#
TensorDictBase is the abstract base class of all tensordicts.
Its subclasses keep their entries in different places:
TensorDictkeeps its tensors in memory. It is the default container.LazyStackedTensorDictis a stack of tensordicts that keeps each tensordict separate, so that their entries can have different shapes. See Stacking and variable-size data.PersistentTensorDictreads and writes its entries in an HDF5 or zarr file. See Large datasets and storage backends.TensorDictParamsis a tensordict and atorch.nn.Module. If a module holds one, the parameters and buffers of the module include its tensors.
MemoryMappedTensor is not a tensordict. It is a tensor whose data
is in a file, and the type of the entries of a memory-mapped tensordict.
TensorCollection is the common base class of the tensordicts and
the subclasses of TensorClass: the type hint for “a tensordict or
a tensorclass”.
Other containers have their own pages: TensorClass, the class
form with typed fields and methods, and non-tensor data
(tensorclass and non-tensor data), TypedTensorDict
(TypedTensorDict), and the key-value stores
(Key-value stores and UCXX). To choose one, see Which container should I use?.
|
The abstract base class of every tensordict. |
|
A dictionary of tensors that share a batch size. |
|
A stack of tensordicts that keeps each input as a separate tensordict. |
|
A tensordict whose entries live in an HDF5 file or a Zarr store. |
|
A TensorDict whose entries are registered as the parameters and buffers of an |
|
A CPU tensor whose storage is a file on disk, or anonymous shared memory when no filename is given. |
|
The abstract base class of tensordicts and tensorclasses. |
Methods by task#
The tables below list the methods of TensorDictBase, with one
page for each method. Every tensordict class and every tensorclass has
these methods. A class page lists only the methods that the class adds,
and the overrides that have their own documentation. When the
tensordict module has a function for the same task, the function is
in the same group.
Construction and conversion#
Create a tensordict from Python, numpy, or pandas data, and convert a
tensordict to other types. Most from_* class methods have a function
of the same name in the tensordict module. To read a tensordict from a
file, see Saving, loading, and file formats. For the arguments of the constructor,
see TensorDict.
Methods
|
Converts a dict, tuple, dataclass, array or file to a tensordict, recursively. |
|
Returns a tensordict built from a dictionary, which can be nested. |
|
Builds a tensordict from a dictionary, with the types of the entries of this tensordict. |
|
Converts a named tuple to a tensordict whose keys are its field names. |
|
Converts a pandas DataFrame to a tensordict, with one entry per column. |
|
Converts a NumPy structured array to a tensordict, with one entry per field. |
|
Converts a tuple to a tensordict with the keys |
|
Converts a pytree to a tensordict that |
|
Converts a dataclass instance to a tensordict, with one entry per field. |
|
Builds a tensordict by stacking a list of dicts or tensor collections. |
|
Allocates a zero-filled tensordict from a schema of shapes and dtypes. |
|
Creates a tensordict with the same value for every key, like |
|
Returns a new tensordict of the same type, with no tensor leaves. |
|
Returns a tensordict of batch size |
|
Returns a tensordict of batch size |
|
Returns a tensordict of batch size |
|
Returns a tensordict whose leaves are copies of |
|
Returns a tensordict of batch size |
Returns a shallow copy of the tensordict: a new structure that holds the same leaves. |
|
|
Returns a copy of the tensordict, of the same type. |
|
Returns a nested dictionary with the entries of the tensordict. |
|
Converts the tensordict to a named tuple, recursively. |
|
Converts the tensordict to a pandas DataFrame, with one column per leaf. |
Rebuilds the pytree that |
|
Converts the tensordict to a NumPy structured array, with one field per entry. |
|
|
Returns a |
|
Returns the elements of the tensordict as nested lists, one level per batch dimension. |
Returns a nested dict of NumPy arrays with the entries of the tensordict. |
|
Returns a tensordict whose leaves are converted with their |
Functions
|
Converts a dict, tuple, dataclass, array or file to a tensordict, recursively. |
|
Converts a dictionary to a tensordict, recursively. |
|
Converts a named tuple to a tensordict whose keys are its field names. |
|
Converts a pandas DataFrame to a tensordict, with one entry per column. |
|
Converts a NumPy structured array to a tensordict, with one entry per field. |
|
Converts a tuple to a tensordict with the keys |
|
Converts a pytree (nested lists, tuples and dicts of tensors) to a tensordict. |
|
Creates a tensordict with the same value for every key, like |
|
Returns a tensordict built from a dict or from keyword arguments. |
Keys, values, and entries#
Read, write, select, and rename the entries of a tensordict. A key is a
string, or a tuple of strings that points to an entry of a nested
tensordict. NestedKey is the type of these keys. For the rules,
see Keys, nested keys, and locking. For values that are not tensors, see
Non-tensor data and unbatched entries.
Methods
|
Returns a view of the keys of the tensordict. |
|
Returns a generator of the values of the tensordict. |
|
Returns a generator of the |
The top-level keys of the tensordict, as a list sorted in alphabetical order. |
|
Returns the value stored at |
|
Returns the entry at |
|
|
Sets an entry, and returns the tensordict. |
|
Copies a value into an existing entry, which keeps its storage, and returns |
|
Writes |
|
Returns the entry at |
|
Writes the entries of a dictionary or another tensordict into the tensordict, and returns |
|
Copies the entries of a dictionary or tensordict into the existing entries, and returns |
|
Copies the entries of a dictionary or tensordict into the existing entries at |
|
Copies the entries of |
|
Copies the tensors of |
|
Returns a shallow copy of the tensordict with some entries replaced. |
|
Removes an entry from the tensordict and returns its value. |
Removes the last top-level entry and returns it as a |
|
|
Deletes an entry of the tensordict in place and returns the tensordict. |
Removes every entry of the tensordict and returns |
|
|
Returns a tensordict that holds only the entries of |
|
Returns a tensordict without the entries of |
|
Splits the tensordict into one tensordict per set of keys, plus one for the remaining entries. |
|
Removes the entries of |
|
Renames an entry of the tensordict in place and returns the tensordict. |
|
Returns a tensordict with several keys renamed at once. |
|
Returns a tensordict with the nested entries moved to the top level, under joined keys. |
|
Returns a tensordict with the keys that hold the separator split into nested keys. |
Adds an empty nested tensordict at |
|
the maximum number of levels of nested tensordicts. |
|
Returns the class of an entry, possibly without reading the entry. |
|
Returns the shape of the entry at |
|
Returns |
|
|
Fills the entry at |
|
Stores |
|
Returns the content of the non-tensor entry at |
Returns |
|
|
Returns a tuple of the |
Returns a tensordict without the non-tensor entries, nested ones included. |
|
Removes the nested tensordicts that have no leaf, in place, and returns |
|
|
Concatenates entries of the tensordict into a new entry, and removes them by default. |
|
Stacks entries of the tensordict into a new entry, and removes them by default. |
|
Concatenates all the leaves of the tensordict into one tensor, with |
|
Packs all the leaves of a tensordict into one tensor of features. |
|
Stacks all the leaves of the tensordict into one tensor, with |
Classes and functions
The type of the keys of a tensordict: a string, or a tuple of strings for a nested entry. |
|
|
The layout of a tensordict packed by |
|
Unravels a nested key into a str or a flat tuple of str. |
|
Unravels every key of a list, as |
|
Unravels a nested key into a flat tuple of str. |
|
Merges the entries of several tensordicts into a new tensordict. |
Indexing, masking, and splitting#
Index every entry along the batch dimensions at once, and split a
tensordict into several tensordicts. td[0] and td[mask] index the
batch dimensions. td["key"] returns an entry: see
Keys, values, and entries. For the rules, see Batch size, shapes, and indexing.
Returns the entry at a key, or the tensordict indexed along its batch dimensions. |
|
Iterates over the first batch dimension of the tensordict. |
|
|
Gathers values along a batch dimension, as |
|
Returns a copy of the tensordict with |
|
Writes |
|
Returns a tensordict with the elements of the batch where |
|
Returns a tensordict with the leaves where |
|
Splits the tensordict along a batch dimension, as |
|
Splits the tensordict into |
|
Splits the tensordict along a batch dimension, as |
Returns a tuple of the tensordicts along a batch dimension, as |
|
|
Returns a tensordict narrowed along a batch dimension, as |
|
Splits a tensordict into groups of rows that share the value of an entry. |
Shape, batch dimensions, and stacking#
Read and change the batch size and the dimension names, reshape and expand a tensordict, and stack or concatenate several tensordicts. These methods change the batch dimensions. The dimensions of each entry after the batch dimensions do not change. See Batch size, shapes, and indexing and Stacking and variable-size data.
Methods
The batch size of the tensordict, as a |
|
The number of batch dimensions, that is |
|
the same as |
|
The number of batch dimensions, the same as |
|
Returns the number of batch dimensions, the same as |
|
Returns the number of batch dimensions, the same as |
|
|
Returns the batch size, or the size of one batch dimension. |
Returns the number of elements in the batch, which is the product of |
|
Returns the size of the first batch dimension, or 0 if the tensordict has no batch dimensions. |
|
The dimension names of the tensordict, one per batch dimension. |
|
|
Names the unnamed batch dimensions, and checks the named ones, in place. |
|
Returns a shallow copy of the tensordict with new names for its batch dimensions. |
|
Renames the batch dimensions of the tensordict in place and returns |
Sets the batch size to the largest leading shape that all the entries share. |
|
|
Returns a tensordict whose batch dimensions are reshaped to |
|
Returns a tensordict with a new batch size, whose leaves are views of the original leaves. |
|
Returns a tensordict whose batch dimensions |
|
Returns a tensordict whose batch dimension |
|
Returns a tensordict with its batch dimensions permuted, as |
|
Returns a tensordict with the batch dimensions |
|
Returns a tensordict with two batch dimensions swapped, as |
|
Returns a tensordict with two batch dimensions swapped, as |
|
Returns a tensordict with batch dimensions moved to new positions, as |
|
Returns a tensordict with batch dimensions moved to new positions, as |
|
Returns a tensordict without the batch dimensions of size 1. |
|
Returns a tensordict with a new batch dimension of size 1 at position |
|
Returns a tensordict whose batch dimensions are expanded to |
|
Returns the tensordict expanded to the shape of |
|
Returns the tensordict broadcast to the batch size |
|
Repeats the batch along each batch dimension, as |
|
Repeats each element of the batch along |
|
Returns a tensordict whose batch is repeated, as |
|
Returns a tensordict with the elements reversed along batch dimensions, as |
Returns a tensordict with the elements reversed along batch dimension 1. |
|
Returns a tensordict with the elements reversed along batch dimension 0. |
|
|
Returns a tensordict with the elements rolled along batch dimensions, as |
|
Returns a tensordict rotated by 90 degrees in the plane of two batch dimensions. |
|
Returns the tensordict padded along its batch dimensions with a constant value. |
Returns the tensordict with at least 1 batch dimension. |
|
Returns the tensordict with at least 2 batch dimensions. |
|
Returns the tensordict with at least 3 batch dimensions. |
|
|
Returns a new tensordict with contiguous leaves. |
|
|
|
Stacks tensordicts into one dense tensordict along a new dimension |
|
Concatenates tensordicts along an existing batch dimension |
|
Stacks tensordicts into a lazy stack. |
|
Stacks tensordicts densely if possible, or into a lazy stack otherwise. |
|
Returns a lazy stack of the slices of the tensordict along |
|
Returns the tensordict with its lazy stacks replaced by dense tensordicts. |
|
Returns a tensordict whose nested-tensor leaves are converted to padded tensors. |
Functions
|
Stacks tensordicts into one dense tensordict along a new dimension |
|
Concatenates tensordicts along an existing batch dimension |
|
Stacks tensordicts into a lazy stack ( |
|
Stacks tensordicts densely if possible, or into a lazy stack otherwise. |
|
Stacks a list of tensordicts, or the tensordicts of a lazy stack, densely if possible. |
|
Pads all tensors in a tensordict along the batch dimensions with a constant value. |
|
Pads tensordicts to a common size along one dimension and stacks them densely. |
|
Splits a padded batch back into its variable-length samples. |
|
Collates a list of samples into a batch, for |
Device, dtype, and memory#
Move a tensordict to a device, cast its entries to another dtype, and control the memory that holds them. See Devices, dtypes, and memory.
Device and memory
|
Returns a tensordict cast to another device, dtype or memory format, or with a new batch size. |
|
Returns the tensordict on the CPU, as |
|
Returns the tensordict on a CUDA device, as |
The device of the tensordict, or |
|
Sets the device of the tensordict to the device of its leaves, if they all share one. |
|
Sets the device of the tensordict and its sub-tensordicts to |
|
Whether the tensordict has a CPU device ( |
|
Whether the tensordict has a CUDA device ( |
|
Whether the tensordict has a meta device ( |
|
|
Returns a tensordict whose tensor leaves are in pinned (page-locked) memory. |
Moves the leaves to shared memory, locks the tensordict, and returns |
|
Whether the tensordict is in shared memory. |
|
|
Marks the tensor leaves of the tensordict as used by |
|
Copies all the leaves into one contiguous storage, and returns a tensordict of views of it. |
Whether the leaves of the tensordict are views into one storage made by |
|
|
Reconstructs a tensordict from a file written by |
|
Returns the number of bytes of the tensor leaves. |
|
Returns the data pointers of the leaves, as a tensordict of integers. |
|
Returns the total number of elements of the tensor leaves. |
|
Returns a tensordict that describes each tensor leaf with a |
Clears the record of the last operation of the tensordict and its nested tensordicts. |
Dtype
The dtype of the leaves, or |
|
|
Returns a tensordict whose tensor leaves are cast to |
|
|
|
Returns a tensordict whose tensor leaves are cast to |
|
Returns a tensordict whose tensor leaves are cast to |
|
Returns a tensordict whose tensor leaves are cast to |
|
Returns a tensordict whose tensor leaves are cast to |
|
Returns a tensordict whose tensor leaves are cast to |
|
Returns a tensordict whose tensor leaves are cast to |
|
Returns a tensordict whose tensor leaves are cast to |
|
Returns a tensordict whose tensor leaves are cast to |
|
Returns a tensordict whose tensor leaves are cast to |
|
Returns a tensordict whose tensor leaves are cast to |
|
Returns a tensordict whose tensor leaves are cast to |
|
Returns a tensordict whose tensor leaves are cast to |
|
Returns a tensordict whose tensor leaves are cast to |
|
Returns a tensordict whose tensor leaves are cast to |
|
Returns a tensordict whose tensor leaves are cast to |
|
Returns a tensordict whose tensor leaves are cast to |
|
Returns a tensordict whose tensor leaves are cast to |
|
Returns a tensordict whose tensor leaves are cast to |
|
Returns a tensordict whose tensor leaves are cast to |
|
Returns a tensordict whose tensor leaves are cast to |
to and the dtype shorthands take a cast_filter predicate
fn(key, value) -> bool that selects the leaves to cast. The predicates
below test the dtype of their last argument, so they serve as a
cast_filter and as a test in apply().
Functions
|
Reconstructs a tensordict from a file written by |
|
Whether the last argument is a floating-point or complex tensor. |
|
Whether the last argument is a floating-point tensor. |
|
Whether the last argument is an integer tensor (signed or unsigned, not bool). |
|
Whether the last argument is a boolean tensor. |
Pointwise math#
Apply arithmetic and elementwise functions to every entry. The Python
operators (+, -, *, /, **, ==, &, …) call
these methods. A method whose name ends with an underscore, such as
add_, writes into the existing tensors: see
In-place operations. For broadcasting and for tensordicts with
different keys, see Arithmetic, apply, reductions, and gradients.
Arithmetic
|
Returns a new tensordict with |
|
Adds |
|
Returns a new tensordict with |
|
Subtracts |
|
Returns a new tensordict with every leaf, scaled by |
|
Returns a new tensordict with every leaf multiplied by |
|
Multiplies every leaf by |
|
Returns a new tensordict with every leaf divided by |
|
Divides every leaf by |
|
Returns a new tensordict with every leaf raised to the power |
|
Raises every leaf to the power |
|
Returns a new tensordict with the remainder of every leaf divided by |
|
Returns a new tensordict that interpolates linearly from every leaf (the start) to |
|
Interpolates linearly from every leaf to |
|
Returns a new tensordict with |
|
Adds |
|
Returns a new tensordict with |
|
Adds |
|
Returns a new tensordict with every leaf clamped into the range |
|
Returns a new tensordict with every leaf clamped to at least |
|
Clamps every leaf to at least |
|
Returns a new tensordict with every leaf clamped to at most |
|
Clamps every leaf to at most |
|
Returns a new tensordict with the element-wise maximum of every leaf and |
|
Replaces every leaf with its element-wise maximum with |
|
Returns a new tensordict with the element-wise minimum of every leaf and |
|
Replaces every leaf with its element-wise minimum with |
Fills every leaf with zeros in place and returns |
Elementwise functions
Returns a new tensordict with |
|
Applies |
|
Returns a new tensordict with |
|
Applies |
|
Returns a new tensordict with |
|
Applies |
|
Returns a new tensordict with |
|
Applies |
|
Returns a new tensordict with |
|
Applies |
|
Returns a new tensordict with |
|
Applies |
|
Returns a new tensordict with |
|
Applies |
|
Returns a new tensordict with |
|
Applies |
|
Returns a new tensordict with |
|
Applies |
|
Returns a new tensordict with |
|
Applies |
|
Returns a new tensordict with |
|
Applies |
|
Returns a new tensordict with |
|
Applies |
|
Returns a new tensordict with |
|
Applies |
|
Returns a new tensordict with |
|
Applies |
|
Returns a new tensordict with |
|
Applies |
|
Returns a new tensordict with |
|
Applies |
|
Returns a new tensordict with |
|
Applies |
|
Returns a new tensordict with |
|
Applies |
|
Returns a new tensordict with |
|
Applies |
|
Returns a new tensordict with |
|
Applies |
|
Returns a new tensordict with |
|
Applies |
|
Returns a new tensordict with |
|
Applies |
|
Returns a new tensordict with |
|
Applies |
|
Returns a new tensordict with |
|
Applies |
|
Returns a new tensordict with |
|
Applies |
|
Returns a new tensordict with |
|
Applies |
|
Returns a new tensordict with |
|
Applies |
|
Returns a new tensordict with |
|
Applies |
|
Returns a new tensordict with |
|
Applies |
|
|
Returns a new tensordict with |
Comparisons and logical operations
|
Compares the tensordict with |
|
Compares the tensordict with |
|
Compares the tensordict with |
|
Compares the tensordict with |
|
Compares the tensordict with |
|
Compares the tensordict with |
|
Computes |
|
Computes |
|
Computes |
Returns |
|
|
Returns a new tensordict with the bitwise AND of every leaf and |
|
Returns a new tensordict with the logical AND of every leaf and |
Returns a new tensordict with |
|
Returns a new tensordict with |
|
Returns a new tensordict with |
|
Returns a new tensordict with |
|
Returns a new tensordict with |
Reductions#
Reduce each entry over all its elements, or along batch dimensions with
dim. With reduce=True, most reductions return one tensor for the
whole tensordict. See Arithmetic, apply, reductions, and gradients.
|
Returns whether every element of every leaf is |
|
Returns whether any element of any leaf is |
|
Returns the maximum of every leaf, over all its elements or along a batch dimension. |
|
Returns the minimum of every leaf, over all its elements or along a batch dimension. |
|
Returns the maximum of every leaf, over all its elements or along a batch dimension. |
|
Returns the minimum of every leaf, over all its elements or along a batch dimension. |
|
Returns the cumulative maximum of every leaf along a batch dimension. |
|
Returns the cumulative minimum of every leaf along a batch dimension. |
|
Returns the mean of every leaf, over all its elements or along batch dimensions. |
|
Returns the mean of the non-NaN elements of every leaf. |
|
Returns the sum of every leaf, over all its elements or along batch dimensions. |
|
Returns the sum of the non-NaN elements of every leaf. |
|
Returns the product of every leaf, over all its elements or along a batch dimension. |
|
Returns the standard deviation of every leaf, over all its elements or along batch dimensions. |
|
Returns the variance of every leaf, over all its elements or along batch dimensions. |
|
Returns the |
|
Returns the log of the summed exponentials of every leaf along batch dimensions. |
|
Returns a tensordict with the norm of each leaf. |
Apply, modules, and gradients#
Apply a function to every entry, and map a function over the batch. Put the parameters of a module in a tensordict and back, and compute gradients. See Arithmetic, apply, reductions, and gradients, Modules: TensorDictModule and tensordict.nn, and torch.compile, torch.func, torch.export, and CUDA graphs.
Methods
|
Applies a callable to every leaf and returns a new tensordict with the results. |
|
Applies a callable to every leaf and writes the results into the existing leaves. |
|
Applies a callable to every leaf and its key, and returns a new tensordict with the results. |
|
Applies a function to chunks of the tensordict in worker processes, and gathers the results. |
|
Applies a function to chunks of the tensordict in worker processes, and yields the results. |
|
Returns a tensordict of the parameters and buffers of a module. |
|
Stacks the parameters and buffers of several modules in one tensordict. |
|
Writes the content of the tensordict into the parameters and buffers of a module, recursively. |
|
Computes the gradients of the leaves of the tensordict that require gradients. |
|
|
|
Sets |
Returns a new tensordict with every tensor leaf detached from the autograd graph. |
|
Detaches every tensor leaf from the autograd graph in place, and returns the tensordict. |
|
A tensordict of the |
|
the same storage, without autograd. |
|
|
Resets the |
Functions
|
Returns a locked tensordict that holds the parameters and buffers of a module. |
|
Stacks the parameters and buffers of several modules in one tensordict. |
|
Returns |
|
Returns |
Saving, loading, and file formats#
Save a tensordict to disk and load it again, and write its entries to
memory-mapped files. Other methods read and write key-value stores, and HDF5,
zarr, CSV, JSON, Parquet, and MDS files, and convert to and from Arrow
tables.
memmap(return_early=True) returns a
TensorDictFuture. See Save and load tensordicts and
Large datasets and storage backends.
Methods
|
Saves the tensordict to disk, and returns a memory-mapped copy of it. |
|
|
|
Saves the tensordict to disk, and returns a memory-mapped copy of it. |
|
Writes every leaf to a memory-mapped file and returns a new tensordict of them. |
|
Writes every leaf to a memory-mapped file in place, and returns |
|
Creates a memory-mapped tensordict with the same keys, shapes and dtypes, without copying the values. |
|
Reloads the tensordict from its |
|
Loads a tensordict that |
|
Loads a saved tensordict into the tensordict, and returns |
|
Adds an uninitialized memory-mapped tensor of shape |
|
Adds a memory-mapped tensor at |
|
Adds a memory-mapped copy of |
Whether the tensordict is memory-mapped. |
|
The directory where a memory-mapped tensordict stores its files, as a |
|
|
Returns the leaves of the tensordict in an |
|
Loads a state dict, as |
|
Opens an HDF5 file as a |
|
Writes the tensordict to a new HDF5 file and returns a tensordict backed by it. |
|
Opens a zarr store as a |
|
Writes the tensordict to a new zarr store and returns a tensordict backed by it. |
|
Reads a CSV file into a tensordict, with one entry per column. |
|
Writes the tensordict to a CSV file, with one column per leaf. |
|
Reads a JSON file of records into a tensordict, with one entry per field. |
|
Writes the tensordict to a JSON file, as an array of records. |
|
Reads a Parquet file into a tensordict, with one entry per column. |
|
Writes the tensordict to a Parquet file, with one column per leaf. |
|
Writes the tensordict to a MosaicML Streaming (MDS) dataset. |
|
Uploads the tensordict to a key-value server and returns a tensordict backed by it. |
Functions
|
Saves a tensordict to disk, and returns a memory-mapped copy of it. |
|
Loads a tensordict that |
|
Writes every leaf to a memory-mapped file and returns a new tensordict of them. |
|
|
|
Appends rows to a tensordict saved in a directory, along its first batch dimension. |
|
Packs a memory-mapped tensordict directory into a single-file zip archive. |
|
Extracts a memory-mapped tensordict archive into a directory. |
|
Whether |
|
Recomputes the CRC-32 checksums of a memory-mapped archive after in-place writes. |
|
Saves a tensordict or a tensorclass to a safetensors file. |
|
Loads a tensordict from a safetensors file. |
|
Opens an HDF5 file as a |
|
Opens a zarr store as a |
|
Reads a CSV file into a tensordict, with one entry per column. |
|
Reads a JSON file of records into a tensordict, with one entry per field. |
|
Reads a Parquet file into a tensordict, with one entry per column. |
|
Converts an Arrow table to a tensordict, with one entry per column, viewing the data where it can. |
|
Converts a tensordict with one batch dimension to an Arrow table, with one column per entry. |
Locking#
A locked tensordict raises an error if you add, replace, rename, or delete
an entry. In-place writes with set_ and update_ still work, also for
non-tensor entries. unlock_ also works in a with block: the
tensordict is locked again at the end of the block. See Keys, nested keys, and locking.
Locks the tensordict and its nested tensordicts, and returns it. |
|
Unlocks the tensordict and its nested tensordicts, and returns it. |
|
|
Distributed and remote#
Send and receive tensordicts between processes with
torch.distributed, run collective operations on them, and stream
them over UCXX (see Key-value stores and UCXX). See Distributed tensordicts.
|
Sends the leaves of the tensordict to another rank, with |
|
Receives the leaves that |
|
Sends the leaves of the tensordict to another rank, one asynchronous transfer per leaf. |
|
Receives the leaves that |
|
Reduces each leaf over all the ranks of the group, in place, so that every rank holds the result. |
|
Reduces each leaf over all the ranks of the group, into the tensordict of rank |
|
Sends the tensordict of rank |
|
Scatters a list of tensordicts from |
|
Returns a list with the tensordict of each rank of the group, on every rank. |
|
Stacks the tensordicts of the other ranks into the tensordict of rank |
|
Sends the structure and the content of the tensordict to a rank that does not know them. |
|
Receives a tensordict that |
|
Sends the tensordict through a UCXX pipe, asynchronously. |
|
Receives a tensordict from a UCXX pipe, into the tensordict if possible, asynchronously. |
Deprecated methods#
These methods still work, and each call raises a DeprecationWarning.
Their pages name the replacement and the release that removes them.
|
Loads a saved tensordict into the tensordict. |
|
Pins the tensor leaves of the tensordict and returns |
|
Returns a tensordict whose tensor leaves are cast to |
|
Returns a tensordict whose tensor leaves are cast to |
|
Returns a tensordict whose tensor leaves are cast to |
|
Returns a tensordict whose tensor leaves are cast to |
Other functions#
Type checks and testing#
Test whether an object is a tensordict, a tensorclass, or non-tensor data, and compare two tensordicts in tests.
|
Returns |
|
Returns |
|
Returns |
|
Asserts that two tensordicts have the same keys and close values. |
|
Asserts that two tensordicts have the same keys and close values. |
Printing#
Choose what the representation of a tensordict shows, and rebuild a tensordict from its printed form. See Printing and display.
|
Sets the options that control what the repr of a tensordict shows. |
Returns the current print options of tensordicts as a dict. |
|
Builds a |
Global settings#
These functions read or change a behavior of tensordict, for the whole
process or inside a with block. set_list_to_stack(False) and
set_capture_non_tensor_stack(True) are deprecated, and their pages say
what to use instead.
|
Returns whether assigning a list to a key stores a |
|
Sets whether assigning a list to a key stores a stack of non-tensor data. |
|
Returns whether stacking identical non-tensor data gives one value instead of a stack. |
Sets whether stacking identical non-tensor data gives one value instead of a stack. |
Deprecated functions#
These functions still work. Their pages name the replacement and the release that removes them.
|
Returns whether |
|
Sets whether |
|
Returns whether the legacy lazy mode is on. |
|
Turns the legacy lazy mode on or off. |
Package constants#
- tensordict.KEY_SEPARATOR#
".", the defaultseparatorofflatten_keys()andunflatten_keys(), for code that joins or splits nested keys the same way, such astuple(name.split(KEY_SEPARATOR)). The names ofstate_dict()also use".", thetorch.nn.Moduleconvention, independently of this constant. Reassigning it has no effect: the defaults are bound when the methods are defined. It can also be imported fromtensordict.utils.
- tensordict.base.NO_DEFAULT#
The sentinel that marks an argument as not given where
Noneis a valid value, such asdefaultinget_non_tensor()ordeviceinnew_zeros(). It can also be imported fromtensordict.tensordict.
- tensordict.__version__#
The version string of the installed tensordict package, or
Noneif it cannot be determined, for example when tensordict is imported from a source tree that is not installed.
See also#
What is TensorDict?, Batch size, shapes, and indexing, Keys, nested keys, and locking, Stacking and variable-size data, Arithmetic, apply, reductions, and gradients, and Devices, dtypes, and memory: the user guide pages that explain these methods.
Save and load tensordicts, Large datasets and storage backends, Load data and share it between processes, and Distributed tensordicts: the how-to guides.
Utilities: the helpers of
tensordict.utils.