Ports the graph tracing to torch.fx

[Jegp]

This PR ports the graph tracing to torch.fx instead of using the homebrewed graph tracer.

TBC...

[chanokin]

this is great work, almost ready!

[chanokin]

Does skipped_nodes means nodes we still need to process?

[Jegp]

It means that we can rewire around them. That is, if we want to skip B in A -> B -> C we can avoid B and go straight from A -> C. This is useful when encountering + because in that case, we simply want circumvent + since, in NIR, two edges into the same node is exactly addition.

[chanokin]

Right, so the only thing on this one is that the variable name is kind-a confusing but this is stricktly personal preference. To me skipped hints to the ignore these nodes path, even if you already have the ignored_nodes set there.

[Jegp]

That's actually a good point. What would be a better name? "bypassed_nodes"? :)

[chanokin]

I was thinking more on something in the lines of "rewiring_nodes" or "postprocessed_nodes"

[Jegp]

What about "bypass"?

[Jegp]

I've added a method to translate NIR graphs to Torch graphs using torch.fx. There are a few questions I would be curious to hear your opinion on:

• We previously used a function to map NIR nodes to Torch nodes. I replaced that with a dictionary, where we can use the keys to check for a mapping. That is Dict[nir.NIRNode, Callable[[nir.NIRNode], torch.nn.Module]]. Downsides are that we have to deprecate the old syntax and the dictionary will be a little annoying to build. We'd have to have individual lambdas/functions plugged into the dictionary. Not hard, just nasty syntax. Is it worth it or should we stay with the old function approach?
• Some modules (like the ones in Norse) use stateful function calls and returns tuples of (output, state). That means, the graph should be able to handle state. By default, all modules will work with states. Even if there are no modules with state, we provide None values for each module, just to take care of the few modules that do. An alternative approach would be to store mutable state inside the modules. But I despise that method because we might forget to reset the state, leave it dangling in memory, etc. etc. Inputs are welcome :-)

[chanokin]

just a tiny observation in the _is_stateful and _construct_fx_graph functions.

[chanokin]

Given that you have an "or" at the end, does this really ensure you have a "state" argument in the function? Just becaues _construct_fx_graph depends on this when you have a stateful function.

Also, how would you enforce this in customizations to the frameworks, for example, if John Doe creates a brand new AwesomeNeuron type for Norse, can you raise an exception if the forward method lacks the "state" argument?

[chanokin]

* We previously used a function to map NIR nodes to Torch nodes. I replaced that with a dictionary, where we can use the keys to check for a mapping. That is `Dict[nir.NIRNode, Callable[[nir.NIRNode], torch.nn.Module]]`. Downsides are that we have to deprecate the old syntax and the dictionary will be a little annoying to build. We'd have to have individual lambdas/functions plugged into the dictionary. Not hard, just nasty syntax. Is it worth it or should we stay with the old function approach?

I remember we talked about dictionaries making things easier to grab lists of available Nodes and to add custom nodes without needing to modify codebases for the packages (be it NIR, Norse or SNNTorch).

I'm not entirely sure what the problem with the pointer to a function is? As long as you have somewhat fixed arguments, it's no biggie in my mind.

All that being said, I think this is ready to fly! Minor adjustments could be dealt with later.

[pabogdan]

It's difficult for me to say at the moment with confidence, but this looks good to me (superficially). Mapping such as these should enable more complex extensions / transforms as well if someone needs to implement those

[pabogdan]

Why is the counter set to 3 here? I would expect a single visit to a note is sufficient to disqualify it from looking at it any more

[Jegp]

3 is the third time around (note that the collections.Counter starts counting with 1). Technically, the worst case for the revisit is N - 1 if you have N modules: they could be listed in the opposite order, interchangeably. Here's an example in a simple graph 1 -> 2 -> 3. If the modules are listed in the order 3, 1, 2, the module 3 would re-queue to wait for module 2. Module 1 would be defined, but module 2 would have to wait again. The queue is now 3, 2, in which case module 3 would re-queue to wait for module 2.

[Jegp]

* We previously used a function to map NIR nodes to Torch nodes. I replaced that with a dictionary, where we can use the keys to check for a mapping. That is `Dict[nir.NIRNode, Callable[[nir.NIRNode], torch.nn.Module]]`. Downsides are that we have to deprecate the old syntax and the dictionary will be a little annoying to build. We'd have to have individual lambdas/functions plugged into the dictionary. Not hard, just nasty syntax. Is it worth it or should we stay with the old function approach?

I remember we talked about dictionaries making things easier to grab lists of available Nodes and to add custom nodes without needing to modify codebases for the packages (be it NIR, Norse or SNNTorch).

I'm not entirely sure what the problem with the pointer to a function is? As long as you have somewhat fixed arguments, it's no biggie in my mind.

All that being said, I think this is ready to fly! Minor adjustments could be dealt with later.

The main problem is that we cannot see whether a module is defined without running the code. Having a set of traceable nodes is necessary to mark them as leaf nodes when processing them with torch.fx. We could let the user specify a set of modules in a separate datastructure, but that seems cumbersome...

Thank you to @pabogdan and @chanokin for the comments! If you're ok with it, I'll merge it for now. We can always improve the code later if something needs changing