reference/generated/tensordict.prototype.symbolic_trace

Run in Google Colab

Colab

Download Notebook

Notebook

View on GitHub

GitHub

symbolic_trace

tensordict.prototype.symbolic_trace(td_module: TensorDictModuleBase) → TDGraphModule

Traces a TensorDictModule or a TensorDictSequential into an FX graph.

This function is a prototype: its API can change without a deprecation period. See FX tracing (prototype) for a guide.

It supports a TensorDictModule that wraps an Module that torch.fx.symbolic_trace() can trace, and a TensorDictSequential of such modules. A module that wraps a plain function or a lambda raises a RuntimeError (“Unsupported type … passed for root!”). Other module types, such as ProbabilisticTensorDictModule, raise a TypeError.

Parameters:

td_module (TensorDictModuleBase) – the module to trace.

Returns:

a TDGraphModule with the same inputs and outputs as td_module.

Examples

>>> import torch
>>> from torch import nn
>>> from tensordict import TensorDict
>>> from tensordict.nn import TensorDictModule, TensorDictSequential
>>> from tensordict.prototype.fx import symbolic_trace
>>> seq = TensorDictSequential(
...     TensorDictModule(nn.Linear(3, 4), in_keys=["x"], out_keys=["hidden"]),
...     TensorDictModule(nn.ReLU(), in_keys=["hidden"], out_keys=["y"]),
... )
>>> graph_module = symbolic_trace(seq)
>>> td = graph_module(TensorDict(x=torch.ones(2, 3), batch_size=[2]))
>>> sorted(td.keys())
['hidden', 'x', 'y']
>>> expected = seq(TensorDict(x=torch.ones(2, 3), batch_size=[2]))["y"]
>>> torch.equal(td["y"], expected)
True