How to add custom operators to export_for_training?

torch.compile
1

I am interested in customizing the new torch.export.export_for_training with additional custom operators to add on to the ATen opset.

The docstring indicates that this may be possible pytorch/torch/export/__init__.py at main · pytorch/pytorch · GitHub
“The traced graph (1) produces normalized operators in the all ATen operator set (as well as any user-specified custom operators)”

I have been unable to locate documentation or examples to show how to add my custom operation to the tracing, but am eager to learn more.

Also apologies if I have miscategorized this to torch.compile. I think my question is also relevant to the new pt2e quantization approach.

2

This tutorial might be a good starter explaining different paths to add custom operators (the first link shows how to add custom ops for torch.compile and torch.export).

1 Like