search

kmheckel / spyx

Spyx: Spiking Neural Networks in JAX
https://spyx.readthedocs.io/en/latest/
MIT License
99 stars 11 forks source link

Investigate Pallas to implement faster neurons? #24

Open kmheckel opened 8 months ago

kmheckel commented 8 months ago

https://jax.readthedocs.io/en/latest/pallas/design.html

Create pallas kernel for LIF neurons and investigate to see if it's faster than just static unrolling the LIF as defined in Spyx.

kmheckel commented 7 months ago

https://docs.taichi-lang.org/

Taichi might be another great option for including enhanced/high performance kernels into the system?

kmheckel commented 7 months ago

A big question though is whether this is worth the extra lift since Spyx with native JAX is already within 5% of SpikingJelly?