TensorDictModuleBase#

class tensordict.nn.TensorDictModuleBase(*args: Any, **kwargs: Any)#

Bases: Module

Base class of the tensordict modules, which read in_keys from a tensordict and write out_keys.

Subclass it to write a module whose forward takes a tensordict and returns a tensordict (tensordict_out = module(tensordict_in)). Set in_keys and out_keys as class attributes or in __init__: TensorDictSequential and the other modules read them to know which entries the module reads and writes. To wrap a module that takes and returns tensors, use TensorDictModule instead. See Modules: TensorDictModule and tensordict.nn.

The base class gives every module select_out_keys() and reset_out_keys(), to choose which outputs a call keeps, and reset_parameters_recursive().

Examples

>>> from tensordict import TensorDict
>>> from tensordict.nn import TensorDictModuleBase
>>> class Mod(TensorDictModuleBase):
...     in_keys = ["a"] # can also be specified during __init__
...     out_keys = ["b", "c"]
...     def forward(self, tensordict):
...         b = tensordict["a"].clone()
...         c = b + 1
...         return tensordict.replace({"b": b, "c": c})
>>> mod = Mod()
>>> td = mod(TensorDict(a=0))
>>> td["b"]
tensor(0)
>>> td["c"]
tensor(1)
>>> _ = mod.select_out_keys("c")
>>> td = mod(TensorDict(a=0))
>>> td["c"]
tensor(1)
>>> "b" in td
False

property in_keys#: The keys that the module reads from the input tensordict.

static is_tdmodule_compatible(module)#

Returns True if module has both an in_keys and an out_keys attribute.

TensorDictSequential wraps the callables that fail this check in a WrapModule.

Parameters:: module (object) – the object to check.

property out_keys#

The keys that the module writes to the output tensordict.

After select_out_keys(), only the selected keys. The full list is in out_keys_source.

property out_keys_source#: All the keys that the module writes, before select_out_keys().

reset_out_keys()#

Undoes select_out_keys(): the module writes all its out_keys again.

Returns:: self.

Examples

>>> import torch
>>> from tensordict import TensorDict
>>> from tensordict.nn import TensorDictModule
>>> mod = TensorDictModule(
...     lambda x, y: (x + 2, y + 2), in_keys=["a", "b"], out_keys=["c", "d"]
... )
>>> _ = mod.select_out_keys("d")
>>> mod.out_keys
['d']
>>> _ = mod.reset_out_keys()
>>> mod.out_keys
['c', 'd']
>>> td = mod(TensorDict(a=torch.zeros(()), b=torch.ones(())))
>>> sorted(td.keys())
['a', 'b', 'c', 'd']

reset_parameters_recursive(parameters: TensorDictBase | None = None) → TensorDictBase | None#

Resets the parameters of every submodule that has a reset_parameters method.

Without parameters, the parameters of the module are reset in place. With parameters, the module is left as it is: the values in parameters are reset in place, as if they were loaded into the module (see to_module()). Use this for parameters kept outside the module, for example in a functional call. Values that share their storage with the module parameters, such as those of from_module(), reset the module too.

Parameters:: parameters (TensorDictBase, optional) – the parameters to reset, with the structure of from_module() and an empty batch size. Defaults to None, which resets the parameters of the module.
Returns:: None without parameters. Otherwise a new tensordict of tensors that share their storage with parameters and require gradients.
Raises:: RuntimeError – if parameters has batch dimensions.

Without parameters, the call warns if no submodule has a reset_parameters method.

Examples

>>> import torch
>>> from tensordict import TensorDict
>>> from tensordict.nn import TensorDictModule
>>> from torch import nn
>>> net = nn.Sequential(nn.Linear(2, 3), nn.ReLU())
>>> old_weight = net[0].weight.clone()
>>> module = TensorDictModule(net, in_keys=["x"], out_keys=["y"])
>>> module.reset_parameters_recursive()
>>> torch.equal(old_weight, net[0].weight)
False

The parameters can also live outside the module:

>>> params = TensorDict.from_module(module).clone()
>>> old_params = params.clone()
>>> _ = module.reset_parameters_recursive(params)
>>> (old_params == params).any()
False
>>> torch.equal(net[0].weight, old_params["module", "0", "weight"])
True

select_out_keys(*out_keys) → TensorDictModuleBase#

Keeps only some of the out_keys in the output of the module, and returns the module.

Use it to drop intermediate outputs, for example those of a TensorDictSequential. The module changes in place: out_keys becomes the selected keys, and out_keys_source keeps the full list. A call still computes all the outputs, and then removes the ones that are not selected. It removes nothing else, so the input entries stay in the output. When you call the module with tensors (see dispatch()), it returns only the selected outputs. reset_out_keys() undoes the selection.

Parameters:: *out_keys (NestedKey) – the keys to keep. Each one must be in out_keys.
Returns:: self.
Raises:: ValueError – if a key is not in out_keys. The error names the missing keys, and is also a RuntimeError.

Examples

>>> import torch
>>> from tensordict import TensorDict
>>> from tensordict.nn import TensorDictModule
>>> mod = TensorDictModule(
...     lambda x, y: (x + 2, y + 2), in_keys=["a", "b"], out_keys=["c", "d"]
... )
>>> _ = mod.select_out_keys("d")
>>> mod.out_keys, mod.out_keys_source
(['d'], ['c', 'd'])
>>> td = mod(TensorDict(a=torch.zeros(()), b=torch.ones(())))
>>> sorted(td.keys())
['a', 'b', 'd']
>>> mod(torch.zeros(()), torch.ones(()))
tensor(3.)
>>> _ = mod.reset_out_keys()
>>> sorted(mod(TensorDict(a=torch.zeros(()), b=torch.ones(()))).keys())
['a', 'b', 'c', 'd']

It works on a TensorDictSequential too:

>>> from tensordict.nn import TensorDictSequential
>>> seq = TensorDictSequential(
...     TensorDictModule(lambda x: x + 1, in_keys=["x"], out_keys=["y"]),
...     TensorDictModule(lambda y: y + 1, in_keys=["y"], out_keys=["z"]),
... )
>>> _ = seq.select_out_keys("z")
>>> sorted(seq(TensorDict(x=torch.zeros(()))).keys())
['x', 'z']