Let me address a few questions so you can help me and I can help you.
Q Can we (will be able to) register a JIT pass that may/will affect the “gradient computation” ?
Registering a pass, as for code optimization, is elegant and you and I can work independently. The HW people can change their minds to suite different HW configurations; I can create versions; and our packages are imported afterwords.  The main concern is that you will need to give us partial control when to call the pass.  This means that I can break a perfectly fine JIT.
In your proposal, you will give us the opportunity to activate “the pass”, which is great. Next question.
Q Can we target layers that affect “the gradient computation” ?
I know quantization of convolutions is HW dependent. This means that I will need space to wiggle in choosing the convolution layer and  what do do with it. In practice, the HW will dictate the computation shape and format: in turn, the computation of the convolution. The computation is still based on basic operations such as aten::conv and addition of the bias, I know the HW people want them separated.
Q is it possible to have a tutorial I can build upon so I can practice on the master code for:


(Any) layer selection into a subgraph.


Where to locate such a phase so that I do the least damage to the JIT optimization process


The documentation is great but I learned more from a post like this: https://zasdfgbnm.github.io/2018/09/20/PyTorch-JIT-Source-Code-Read-Note/ 11