neuromorphs / NIRTorch

PyTorch helper module to translate to and from NIR
https://neuroir.org/docs
BSD 3-Clause "New" or "Revised" License
7 stars 3 forks source link

Allow stateful modules in the graph executor #9

Closed Jegp closed 12 months ago

Jegp commented 1 year ago

The graph executor currently doesn't allow modules to have state. If we have a tensor where the outer dimension is time, then we are forced to run all the frames through each node, like so:

tensor = ... # Time x Batch x ...
module(tensor)

We might want to support iterative application, which is for instance necessary for real-time modules:

tensor = ... # Time x Batch x ...
state = None
for frame in tensor:
  out, new_state = module(frame)

This is the (functional) model Norse is following, so we would have to make sure this also works for libraries with mutable state (where the state variable is not explicitly passed).