reference/generated/tensordict.nn.TensorClassModuleBase

Run in Google Colab

Colab

Download Notebook

Notebook

View on GitHub

GitHub

TensorClassModuleBase

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

Bases: Generic[InputClass, OutputClass], ABC, Module

Base class of the modules that take and return a tensorclass.

Subclass TensorClassModuleBase[InputClass, OutputClass], where each type is a TensorClass subclass or torch.Tensor, and implement forward(). A type checker then knows the type of the input and output, and can report a wrong field name. The class reads both types from the subscript: a subclass without it has no input_type and output_type, and its as_td_module() raises a ValueError.

When both types are tensorclasses, as_td_module() returns a module that takes and returns a tensordict, for example to put the module in a TensorDictSequential. See TensorClass: a dataclass of tensors.

Variables:

• input_type (type) – the input type, InputClass.

• output_type (type) – the output type, OutputClass.

Examples

>>> from tensordict.tensorclass import TensorClass
>>> from tensordict.nn import TensorClassModuleBase
>>> import torch
>>>
>>> class InputTC(TensorClass):
...     a: torch.Tensor
...     b: torch.Tensor
...
>>> class OutputTC(TensorClass):
...     result: torch.Tensor
...
>>> class AddModule(TensorClassModuleBase[InputTC, OutputTC]):
...     def forward(self, x: InputTC) -> OutputTC:
...         return OutputTC(
...             result=x.a + x.b,
...             batch_size=x.batch_size
...         )
...
>>> module = AddModule()
>>> input_tc = InputTC(a=torch.tensor([1.0]), b=torch.tensor([2.0]), batch_size=[1])
>>> module(input_tc).result
tensor([3.])

as_td_module() → TensorClassModuleWrapper

Returns a TensorClassModuleWrapper of the module, which takes and returns a tensordict.

The wrapper is a TensorDictModuleBase. Its in_keys and out_keys are the fields of the input and output classes. It returns a new TensorDict with only the output fields (see TensorClassModuleWrapper).

Raises:

ValueError – if input_type or output_type is not a tensorclass, or if the subclass did not set each of them to a class with the subscript.

Examples

>>> from tensordict import TensorDict
>>> from tensordict.tensorclass import TensorClass
>>> from tensordict.nn import TensorClassModuleBase
>>> import torch
>>>
>>> class InputTC(TensorClass):
...     x: torch.Tensor
...
>>> class OutputTC(TensorClass):
...     y: torch.Tensor
...
>>> class MyModule(TensorClassModuleBase[InputTC, OutputTC]):
...     def forward(self, input: InputTC) -> OutputTC:
...         return OutputTC(y=input.x * 2, batch_size=input.batch_size)
...
>>> module = MyModule()
>>> td_module = module.as_td_module()
>>> td = TensorDict({"x": torch.ones(3)}, batch_size=[3])
>>> td_module(td)["y"]
tensor([2., 2., 2.])

abstract forward(x: InputClass) → OutputClass

Computes the output of the module. Subclasses implement it.

Parameters:

x (InputClass) – the input.

Returns:

An instance of OutputClass.