thunder.dynamo.ThunderCompiler

class thunder.dynamo.ThunderCompiler(**thunder_options)

Bases: object

A class that compiles a torch.fx.GraphModule to a thunder.ThunderModule. This class is meant to be used as a backend for the torch.compile() function.

Keyword Arguments:

thunder_options – a dictionary of options to pass to thunder.jit().

Example

>>> import torch
>>> from thunder.dynamo import ThunderCompiler
>>> backend = ThunderCompiler()
>>> x = torch.ones(2, requires_grad=True)
>>> @torch.compile(backend=backend)
... def func(x):
...     x = torch.sin(x)
...     if x.sum() > 0:
...         return x + 1
...     else:
...         return x - 1
>>> out = func(x)

__init__(**thunder_options)

A class that compiles a torch.fx.GraphModule to a thunder.ThunderModule. This class is meant to be used as a backend for the torch.compile() function.

Keyword Arguments:

thunder_options – a dictionary of options to pass to thunder.jit().

Example

>>> import torch
>>> from thunder.dynamo import ThunderCompiler
>>> backend = ThunderCompiler()
>>> x = torch.ones(2, requires_grad=True)
>>> @torch.compile(backend=backend)
... def func(x):
...     x = torch.sin(x)
...     if x.sum() > 0:
...         return x + 1
...     else:
...         return x - 1
>>> out = func(x)

Methods

__init__(**thunder_options)	A class that compiles a torch.fx.GraphModule to a thunder.ThunderModule.
save_reproducer_to_folder(reproducer_folder)	Save the reproducer script for the GraphModule executed by Thunder to the specified reproducer_folder.

__call__(gm, sample_args)

Call self as a function.

Parameters:

• gm (GraphModule) –

• sample_args (list[torch.SymInt, torch.Tensor]) –

save_reproducer_to_folder(reproducer_folder, use_pytest_benchmark=False, serialize_inputs=False)

Save the reproducer script for the GraphModule executed by Thunder to the specified reproducer_folder. Each saved script is named as “graph[graph_id]_thunder_[module_id]”, where:

• graph_id indexes the graph generated by Dynamo, which is then passed to Thunder.

• module_id indexes the submodule split by the thunder.dynamo.utils._splitter().

Parameters:

• reproducer_folder (str | PathLike) – The folder where the reproducer code will be written. Can be specified as an absolute or relative path.

• use_pytest_benchmark (bool) – Determines the type of script to create. When False, create a reproducer script. Otherwise, creats a benchmark script to compare the reproducer’s performance with other backends, including Torch eager, torch.compile.

• reproducer_folder (str | os.PathLike) –

• use_pytest_benchmark (bool) –