Customizing generated code?

torch.compile FX (Functional Transformations)
1

Hi all,

In the context of this issue ([FX] torch.fx generate code that is not torch.jit able · Issue #53517 · pytorch/pytorch · GitHub which is since resolved, though not yet in the stable release) I was curious" is it possible to customize the generated code for a given node?

Thanks!

2

Hello @Linux-cpp-lisp,

This is a good question. We didn’t foresee the desire to do this, so the current APIs aren’t factored in a clean enough way to easily support something like this. (You can override fx.Graph.python_code, but it would involve a lot of copy-paste, which might be OK for some situations). This seems like a good idea, though, and I’ve filed a feature request: [FX] Allow customizing code generation for certain Nodes · Issue #53732 · pytorch/pytorch · GitHub

For issue #53517 in particular, the TorchScript fix is being rolled up into the 1.8.1 patch release, with release date to be announced

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