braincore.surrogate.piecewise_quadratic#
- braincore.surrogate.piecewise_quadratic(x, alpha=1.0)[source]#
Judge spiking state with a piecewise quadratic function [1] [2] [3] [4] [5].
If origin=False, computes the forward function:
\[\begin{split}g(x) = \begin{cases} 1, & x \geq 0 \\ 0, & x < 0 \\ \end{cases}\end{split}\]If origin=True, computes the original function:
\[\begin{split}g(x) = \begin{cases} 0, & x < -\frac{1}{\alpha} \\ -\frac{1}{2}\alpha^2|x|x + \alpha x + \frac{1}{2}, & |x| \leq \frac{1}{\alpha} \\ 1, & x > \frac{1}{\alpha} \\ \end{cases}\end{split}\]Backward function:
\[\begin{split}g'(x) = \begin{cases} 0, & |x| > \frac{1}{\alpha} \\ -\alpha^2|x|+\alpha, & |x| \leq \frac{1}{\alpha} \end{cases}\end{split}\]>>> import braincore as bc >>> import matplotlib.pyplot as plt >>> xs = jax.numpy.linspace(-3, 3, 1000) >>> for alpha in [0.5, 1., 2., 4.]: >>> grads = bc.vector_grad(bc.surrogate.piecewise_quadratic)(xs, alpha) >>> plt.plot(xs, grads, label=r'$\alpha$=' + str(alpha)) >>> plt.legend() >>> plt.show()
- Parameters:
x (jax.Array, Array) – The input data.
alpha (float) – Parameter to control smoothness of gradient
- Returns:
out – The spiking state.
- Return type:
jax.Array
References