reference/nn

Run in Google Colab

Colab

Download Notebook

Notebook

View on GitHub

GitHub

tensordict.nn

This page lists the modules of tensordict.nn, and the probabilistic modules, distributions, and helpers that go with them. A module reads the entries named in its in_keys from a tensordict, and writes the entries named in its out_keys into it. For a guide with examples, see Modules: TensorDictModule and tensordict.nn.

The page of a module class lists its methods and the methods of its tensordict base classes. For the methods of torch.nn.Module, see the PyTorch documentation.

Modules

• TensorDictModuleBase is the base class of the tensordict modules.

• TensorDictModule wraps a torch.nn.Module that takes and returns tensors, and gives it in_keys and out_keys.

• TensorDictSequential runs modules one after the other on the same tensordict.

• WrapModule wraps a function that takes and returns a tensordict.

• TensorDictModuleWrapper wraps a tensordict module and behaves like it, except for the methods that a subclass overrides.

• as_tensordict_module is a decorator that turns a function into a TensorDictModule.

• CudaGraphModule captures a module or a function in a CUDA graph and replays it. See torch.compile, torch.func, torch.export, and CUDA graphs.

TensorDictModuleBase(*args, **kwargs)	Base class of the tensordict modules, which read in_keys from a tensordict and write out_keys.
TensorDictModule(module, in_keys, out_keys, *)	Wraps a callable so that it reads its inputs from a tensordict and writes its outputs to it.
TensorDictSequential(*modules[, ...])	Runs TensorDictModules one after the other on the same tensordict.
WrapModule(func, *[, inplace, in_keys, out_keys])	Wraps a function that takes and returns a tensordict as a tensordict module.
TensorDictModuleWrapper(td_module)	Wraps a tensordict module and behaves like it, except for the methods that a subclass overrides.

as_tensordict_module(*, in_keys, out_keys)	Turns a function or a method into a tensordict module that reads in_keys and writes out_keys.
CudaGraphModule(module[, warmup, in_keys, ...])	Records a function, or a TensorDictModule, as a CUDA graph and replays it.

Modules with tensorclass inputs and outputs

A TensorClassModuleBase subclass declares the tensorclass types of its input and output, so that a type checker can report a wrong field name or type. TensorClassModuleBase.as_td_module() returns a TensorClassModuleWrapper, which reads and writes a tensordict, so that you can put the module in a TensorDictSequential.

TensorClassModuleBase(*args, **kwargs)	Base class of the modules that take and return a tensorclass.
TensorClassModuleWrapper(module)	Wraps a TensorClassModuleBase so that it takes and returns a tensordict.

Probabilistic modules

A ProbabilisticTensorDictModule builds a distribution from entries of the tensordict, and writes a sample, and optionally its log-probability, to the tensordict. InteractionType and set_interaction_type choose how the module gets its output from the distribution: a random sample, the mode, the median, the mean, or another deterministic value. set_composite_lp_aggregate chooses whether a CompositeDistribution aggregates the log-probabilities and entropies of its parts in one tensor.

ProbabilisticTensorDictModule(in_keys, ...)	A module that builds a distribution from entries of a tensordict and writes a value from it.
ProbabilisticTensorDictSequential(*modules)	A TensorDictSequential with probabilistic modules, which can return their distribution.

InteractionType(value[, names, module, ...])	The rule that a probabilistic module uses to get a value from a distribution.
set_interaction_type([type, thread_local])	Sets the interaction type of all probabilistic modules inside a block of code, or inside a decorated function.
set_composite_lp_aggregate([mode, thread_local])	Sets whether a CompositeDistribution sums its log-probabilities in one tensor.

composite_lp_aggregate([nowarn])	Returns whether a CompositeDistribution sums its log-probabilities in one tensor.
interaction_type()	Returns the interaction type set by set_interaction_type, or None outside of it.

Ensembles

EnsembleModule makes num_copies copies of a module, each with its own newly initialized parameters, and calls all of them on the input with torch.vmap(). The output has a new leading batch dimension of size num_copies. To keep parameters as a tensordict on a module, see TensorDictParams on tensordict: containers and functions.

functional_vmap() calls a module once per set of parameters in a stacked tensordict, such as one built with from_modules(), and stacks the outputs. It runs the sets with torch.vmap(), or one at a time for modules that torch.vmap() cannot handle.

EnsembleModule(module, num_copies[, ...])

Runs num_copies copies of a module, each with its own parameters, on the same input.

functional_vmap(module, *, params[, args, ...])

Calls module once per set of parameters in params and stacks the outputs.

Compiling modules

The modules of tensordict.nn work with torch.compile(). For example, you can compile a TensorDictSequential. For the limits and their workarounds, see torch.compile, torch.func, torch.export, and CUDA graphs. To capture a module in CUDA graphs, use CudaGraphModule.

Distributions

These distributions and modules build the distributions of probabilistic modules. NormalParamExtractor and AddStateIndependentNormalScale turn the output of a network into the location and scale of a normal distribution. CompositeDistribution groups several distributions, one for each entry of a tensordict.

AddStateIndependentNormalScale([...])	A module that returns its input as a location, with a learned scale that does not depend on the input.
CompositeDistribution(params, ...[, ...])	A distribution made of several named distributions, which samples and scores tensordicts.
Delta(param[, atol, rtol, batch_shape, ...])	A distribution with all its mass on one value, param.
NormalParamExtractor([scale_mapping, scale_lb])	A module without parameters that splits a tensor into the location and scale of a normal distribution.
OneHotCategorical([logits, probs])	A categorical distribution whose values are one-hot tensors.
TruncatedNormal(loc, scale, a, b[, ...])	A normal distribution of mean loc and standard deviation scale, restricted to [a, b].
truncated_normal.TruncatedStandardNormal(a, b)	A standard normal distribution restricted to the interval [a, b].

Utilities

• dispatch lets you call a module with keyword arguments instead of a tensordict.

• set_skip_existing and skip_existing() make a module skip its computation when its output entries are already in the tensordict.

• mappings(), add_custom_mapping(), and biased_softplus map the output of a network to positive values, for example the scale of a distribution. inv_softplus() is the inverse of the softplus function, and expln() is the positive mapping of the "expln" key of mappings().

• rand_one_hot() samples a one-hot vector from logits.

tensordict.nn.make_tensordict is tensordict.make_tensordict().

dispatch([separator, source, dest, ...])	Lets a function that takes a tensordict be called with tensors as arguments.
set_skip_existing([mode, in_key_attr, ...])	Makes modules skip their computation when their outputs are already in the tensordict.
biased_softplus(bias[, min_val])	A softplus shifted so that it returns bias for a zero input and stays above min_val.

skip_existing()	Returns the mode set by set_skip_existing: False (the default), True, or a list of keys.
mappings(key)	Returns the function named by key, which maps a tensor elementwise, for example to positive values.
add_custom_mapping(name, mapping)	Registers mapping under name, so that mappings() returns it.
inv_softplus(bias)	Returns the input of the softplus function that gives bias: log(exp(bias) - 1).
utils.expln(x)	Maps real numbers to positive numbers: exp(x) for x <= 0, log1p(x) + 1 for x > 0.
rand_one_hot(values[, do_softmax])	Draws a one-hot sample along the last dimension of values.

FX tracing (prototype)

symbolic_trace() traces a tensordict module with torch.fx and returns a TDGraphModule. Both are prototypes. See FX tracing (prototype).

symbolic_trace(td_module)

Traces a TensorDictModule or a TensorDictSequential into an FX graph.

TDGraphModule(graph_module, out_keys)

The module that symbolic_trace() returns: an FX graph that reads and writes a tensordict.

See also

• Modules: TensorDictModule and tensordict.nn: how a module reads and writes a tensordict, with examples.

• torch.compile, torch.func, torch.export, and CUDA graphs: compile modules and capture them in CUDA graphs.

• The tutorials Building models with TensorDictModule, Functional calls and ensembles with TensorDict, and Exporting tensordict modules.