Heads up on Spyx API change

[kmheckel]

Hi!

Just wanted to let you know that the Spyx performance benchmark will throw an error for new versions of Spyx as Axon was deprecated. The correct and stable API is now to just specify activation=spyx.axn.arctan(). Additionally, hk.static_unroll() can be used in place of hk.dynamic_unroll() with max int as the number of steps, and the spyx.nn.LIF() cores do not need to be individually wrapped by hk.DeepRNN.

Thought this would be helpful to pass along if/when you have to rerun the performance benchmark to save you debugging time. I know how tedious performance benchmarks can be when you have to try learning 5 different ways to accomplish the same task!

Feel free to ping me with any other Spyx related questions too by the way, always happy to help!

slax/perf_tests/benchmark_jax.py

Line 87 in 65ccf5d

core = hk.DeepRNN([snn.LIF((n_neurons,),beta=0.5, activation=spyx.axn.Axon(spyx.axn.arctan())),])