As you mentioned, autograd just does the vector jacobian product.
In this case, you multiply the 2d matrix of gradients wrt the output (grad_output) again a 4D Jacobian matrix for that op.
The nice thing is that the op is element-wise and so the 4D Jacibian is actually diagonal and what this vector jacobian product boils down to is multiply each element of the input 2D matrix with each element in the diagonal of the Jacobian.
Here, (1. - z*z) is actually the “diagonal” of the jacobian and the result is computed by doing an element-wise product with the gradient flowing back to get the gradient wrt the input.