reference/generated/tensordict.prototype.fx.TDGraphModule

Run in Google Colab

Colab

Download Notebook

Notebook

View on GitHub

GitHub

TDGraphModule

class tensordict.prototype.fx.TDGraphModule(graph_module: GraphModule, out_keys: list[NestedKey])

Bases: Module

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

A call reads the in_keys of the traced module from the input tensordict, runs the graph, and writes the outputs at out_keys, into tensordict_out if given, else into the input. The graph and code attributes are those of the underlying torch.fx.GraphModule.

This class is a prototype: its API can change without a deprecation period. See symbolic_trace() for an example.

Parameters:

• graph_module (GraphModule) – the traced graph.

• out_keys (list of NestedKey) – the keys where the outputs are written.

forward(tensordict: TensorDictBase, tensordict_out: TensorDictBase | None = None, **kwargs) → TensorDictBase

Runs the graph on tensordict and writes the outputs at out_keys.

Parameters:

• tensordict (TensorDictBase) – the input tensordict.

• tensordict_out (TensorDictBase, optional) – the tensordict to write the outputs into. Defaults to None, which writes them into tensordict.

• **kwargs – passed to the graph.

Returns:

The tensordict that holds the outputs.