braincore.surrogate.leaky_relu

braincore.surrogate.leaky_relu(x, alpha=0.1, beta=1.0)

Judge spiking state with the Leaky ReLU function.

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}\begin{split}g(x) = \begin{cases} \beta \cdot x, & x \geq 0 \\ \alpha \cdot x, & x < 0 \\ \end{cases}\end{split}\end{split}\]

Backward function:

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

>>> import braincore as bc
>>> import matplotlib.pyplot as plt
>>> xs = jax.numpy.linspace(-3, 3, 1000)
>>> grads = bc.vector_grad(bc.surrogate.leaky_relu)(xs, 0., 1.)
>>> plt.plot(xs, grads, label=r'$\alpha=0., \beta=1.$')
>>> plt.legend()
>>> plt.show()

(Source code)

Parameters:

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

• alpha (float) – The parameter to control the gradient when \(x < 0\).

• beta (float) – The parameter to control the gradient when \(x >= 0\).

Returns:

out – The spiking state.

Return type:

jax.Array