jukiewiczm:

For example (assuming my data is of shape (n, 10), e.g. n observations 10 features each) will backprop work just fine with such definition of forward function?


Yes, it will work just fine.



 jukiewiczm:

Are there any particular limitations with regard to forward?


Not really. The operations can even be non-differentiable if you don’t want gradients. You should always return some variable through.