@leowalkling,  could you elaborate more on what you want to replace here with a example? In case it’s helpful, in JIT we currently put ops to 2 categories:

Ops have autodiff formulas in autodiff.cpp: their backward graph is replaced by https://github.com/pytorch/pytorch/blob/0cb24098c74f8ebed81ec08b83bf6cb5ab3903f5/torch/csrc/jit/graph_executor.cpp#L76 16

Ops don’t have autodiff formulas: their backwards are handled by eager mode autograd directly.
Please let us know the context of problem so that we can help more on this.

Thanks,
Ailing