Major Revamp:
- Replaced spyx.axn.Axon with spyx.axn.custom
- All preimplemented surrogate gradients are now standalone activation functions, no need to wrap them with another function
- spyx.axn.ActivityRegularization was moved to spyx.nn because it is a hk.Module and therefore a layer.
- Loss and accuracy functions converted to higher order funcs which return the version used in training loops. Time axis arg added.
- Notebooks in the docs were updated to reflect syntax changes.
- time constants are now constrained via jnp.clip which is cleaner and more efficient.
- Fixed bug in data shuffling and optimized it to remove an unnecessary permutation call.