Customizing generated code?

torch.compile FX (Functional Transformations)
3

Hi @James_Reed:

That makes sense. I can think of a few cases where this would be useful:

  • Circumventing compatibility issues (like above)
  • Code style
  • Inserting comments into generated code
  • Inserting things that cannot be represented as typical Nodes, such as profiling blocks with with torch.autograd.profiler.record_function, because they aren’t in the flow of data
  • Inserting simple control flow through custom nodes (i.e. a node that outputs
if my_input is None:
    my_output = torch.zeros(some_shape_node)
else:
    my_output = my_input

Conceptually this is a node — it has two arguments, one a Node and one a shape or a node representing a shape, and always one output. To do this in FX currently, as I understand it, this operation would have to be wrapped in a seperate function that you mark as a leaf node by making a custom Tracer. Doable, but not exactly convenient or modular.

2 Likes