This is very interesting!
I have a somewhat similar question along these lines.  If anyone feels I should address this in a different issue, I can create a new one.
Is there a way to “intercept” the backward pass on-the-fly, as it runs in a normal fashion?  Specifically for pytorch/xla package, but I know ptxla doesn’t have this, so maybe it’s in autograd?  The reason being, i would very much like to be able to tell which forward op each backward op is associated with (i.e. this subgraph is, for the most part, the bwd pass of the dropout layer, for instance) and this be able to keep track of this in some sidebanded fashion to be processed during graph compile?