Ctrl+K

braincore.surrogate.relu_grad

Contents

braincore.surrogate.relu_grad#

braincore.surrogate.relu_grad(x, alpha=0.3, width=1.0)[source]#

Spike function with the ReLU gradient function [1].

The forward function:

\[\begin{split}g(x) = \begin{cases} 1, & x \geq 0 \\ 0, & x < 0 \\ \end{cases}\end{split}\]

Backward function:

\[g'(x) = \text{ReLU}(\alpha * (\mathrm{width}-|x|))\]
>>> import braincore as bc
>>> import matplotlib.pyplot as plt
>>> xs = jax.numpy.linspace(-3, 3, 1000)
>>> for s in [0.5, 1.]:
>>>   for w in [1, 2.]:
>>>     grads = bc.vector_grad(bc.surrogate.relu_grad)(xs, s, w)
>>>     plt.plot(xs, grads, label=r'$\alpha=$' + f'{s}, width={w}')
>>> plt.legend()
>>> plt.show()

(Source code)

Parameters:

  • x (jax.Array, Array) – The input data.

  • alpha (float) – The parameter to control the gradient.

  • width (float) – The parameter to control the width of the gradient.

Returns:

out – The spiking state.

Return type:

jax.Array

References

Contents