Should I sum over one of these batch dimensions?

Not really. If you function is actually independent across batches, then all the non “diagonal” elements will be 0s.
But if you have batchnorm or something like that, they won’t be 0 and you may not want to use them?
If you’re happy with summing, then you can get that result faster by expanding your weights to have a batch size and compute the gradients wrt to these expanded weights.