Add Spiking-RWKV layer to support SpikeGPT architectures

kmheckel / spyx

Spyx: Spiking Neural Networks in JAX

https://spyx.readthedocs.io/en/latest/

MIT License

84 stars 10 forks source link

Add Spiking-RWKV layer to support SpikeGPT architectures #9

Open kmheckel opened 10 months ago

kmheckel commented 10 months ago

It'd be great to provide some template layers within the library for more complex SNN architectures.

[ ] Spiking RWKV layer for SpikeGPT
[ ] Spiking VGG block?
[ ] Spiking mobilenet blocks?
[ ] Spiking Retention?

pengzhouzp commented 9 months ago

This is interesting!

kmheckel commented 9 months ago

Thanks! I'd be happy to include a SpikeGPT/SpikingRWKV block, I just haven't had the chance yet to really dig into the architectures and try implementing it. A numpy or torch implementation of the block would be all I need to translate it into an equivalent JAX module.

pengzhouzp commented 9 months ago

In another RWKV/SpikeGPT project the numpy version of RWKV was done but may not including the spiking block. Would you like to discuss this over email?

kmheckel commented 9 months ago

I'd prefer to discuss things either here or on the Open Neuromorphic discord channel so that way other's could follow along if that works for you!