i.e. extend #199 to all activations, now that Zygote implements the vmap adjoint.
Let's do proper benchmarking when doing this.

@AStupidBear