The second derivative of d^2L / dWo da  is really a tuple of Jacobian matrices, where the ith entry of that tuple is a Jacobian matrix with height of p_i.numel() (where p_i is the ith parameter) and width of a.numel().
What you are actually computing when you compute 'grad is the vjp, i.e., v^T J, where v is the quantity (a tuple of tensors) you pass in as grad_out to .grad(), and J can be thought of as all the Jacobians described above concatenated. Even though you may need to pass in 2-D tensors, the Jacobian really doesn’t care about how they are organized into the shape of a tensor - all that matters is that the number of elements in the tensor becomes the size of row/col of the Jacobian.
So there is really no correct or wrong v here to pass to .grad, whatever you pass as v to .grad just determines the coefficients in the linear combination of the rows of your Jacobian. That way when you pass into grad_out, A’s grad will always be a vector with size like A.
If you really need the entire Jacobian, and not a linear combination of its rows, consider using:
https://pytorch.org/docs/stable/autograd.html?highlight=jacobian#torch.autograd.functional.jacobian
Otherwise, if you need a quantity corresponding to the size of tensor A, and are indifferent to any of the dimensions in each of your parameters, you can just pass all one’s tensors corresponding to the sizes to each of the Tensors you have as parameters.