Implement a backward for Heaviside

Context

Since the Heaviside function is non-differentiable, the classical approach is to use a surrogate gradient for the backward pass. The most common gradient function used is the shifted arctan function :

For the forward pass, th :

S[t] = Heaviside(U[t] - U_{\rm thr}) \approx \frac{1}{π}\text{arctan}(\pi U) \\

Thus for the backward :

\frac{\partial \tilde{S}}{\partial U} = \frac{1}{(1+(U \pi)^2)}