In the documentation, JAX provides vectorization. However, aren't JAX operations already vectorized? For example, to add two vectors, I thought that the element-wise additions were vectorized internally already. My guess is that vectorization is useful when: it's hard for us to add a dimension for broadcasting, so we resort to a more explicit vectorization. EDIT: for example, instead of vectorizing convolution2d with different kernels, I simply stack the kernels, copy and stack the channel, then perform the convolution2d with this stack of kernels. Continue reading...
