CodeGen

class torch.fx.graph.CodeGen

Warning

This API is experimental and is NOT backward-compatible.

additional_globals()

If your codegen uses extra global values, add tuples of (identifier,reference to the value) here. For example, return [‘List’, typing.List] if you need List in the global context.

Return type:

list[tuple[str, Any]]

gen_fn_def(free_vars, maybe_return_annotation, *, expanded_def=False)

Given the free variables and a return annotation, generates the beginning of the FX function. By default, gen_fn_def([‘a’, ‘b’], ‘’) == ‘def {self._func_name}(a, b):’

Return type:

str

generate_output(output_args, *, descs=None, repr_fn=None)

Given the output arguments, generates the return statement of the FX function. Note: The returned statement should not be indented.

Return type:

str

process_inputs(*args)

Transforms the inputs so that the graph can take them as arguments, as non-default codegen may result in the inputs to the function being different from the inputs to the graph.

If the graph was directly runnable, this invariant should hold true f.graph.process_outputs(f.graph(*f.graph.process_inputs(*inputs))) == f(*inputs)

Return type:

Any

process_outputs(outputs)

Transforms the outputs of the graph to be identical to the codegen.

See process_inputs for more details.

Return type:

Any