use vmap for all activations #220
Description
i.e. extend #199 to all activations, now that Zygote implements the vmap adjoint.
Let's do proper benchmarking when doing this.
{{ message }}
i.e. extend #199 to all activations, now that Zygote implements the vmap adjoint.
Let's do proper benchmarking when doing this.