Torch.fx modify a copy of the Graph

torch.compile FX (Functional Transformations)
1

Hi,

Let me explain the problem in its simplest form.
I have a rule for block replacement (block1 → block2), and given a model I would like to generate all possible variations applying this rule.

So what I would like to do is I find all occurences of block1 in the model graph and keep some sort of
reference to the nodes, then I iterate over the binary selection block1/block2 of combinations available and modify the graph accordingly (keep block1 occurence vs delete block1 occurence and insert block2).
The problem is that the graph is modified in place, and I am using memory references to nodes, so after the first iteration I deleted an occurence of block1 I won’t be able to access the node again.
It would be nice if I could modify a copy of the graph, but if I make a copy I loose track of the occurences in the new graph (because of the use of references).

Anyone has a better idea for this?

2

Sorry, I don’t quite understand what your issue is. In particular, I don’t understand the problem here:

so after the first iteration I deleted an occurence of block1 I won’t be able to access the node again.

Could you elaborate a bit more?

3

Sorry I mad it really complex, it is a reference problem.
I am using object reference to reference the nodes composing the block I want to replace, however because modifications are done in place, those references won’t be valid if I wan’t to reuse them on a copy of the model to replace by another block.
I think using node name should be fine for referencing, as long as each node have a unique name, do they?