reference/generated/tensordict.nn.as_tensordict_module

Run in Google Colab

Colab

Download Notebook

Notebook

View on GitHub

GitHub

as_tensordict_module

class tensordict.nn.as_tensordict_module(*, in_keys: List[NestedKey] | NestedKey, out_keys: List[NestedKey] | NestedKey)

Bases: object

Turns a function or a method into a tensordict module that reads in_keys and writes out_keys.

The decorated function is wrapped in a TensorDictModule, built at the first call. Called with a tensordict, it returns the tensordict with the outputs written at out_keys. Called with tensors (see dispatch()), it returns the output tensors. On a method, one module serves all the instances of the class, and the instance is passed as self at each call.

Keyword Arguments:

• in_keys (NestedKey or list of NestedKey) – the keys of the entries passed to the function, in the order of its arguments.

• out_keys (NestedKey or list of NestedKey) – the keys where the outputs of the function are written, one key per output.

Examples

>>> import torch
>>> from tensordict import TensorDict
>>> from tensordict.nn import as_tensordict_module
>>> @as_tensordict_module(in_keys="c", out_keys="d")
... def my_function(c):
...     return c + 1
>>> my_function(TensorDict(c=0))["d"]
tensor(1)
>>> my_function(c=torch.zeros(2))
tensor([1., 1.])

It works on methods too:

>>> class MyClass:
...     @as_tensordict_module(in_keys="c", out_keys="d")
...     def my_method(self, c):
...         return c + 1
>>> MyClass().my_method(TensorDict(c=0))["d"]
tensor(1)