[dattri.algorithm] Fix unsqueeze type artifact for vmap usage

[tingwl0122]

For most of the algorithms we implemented, we use some specific dimension manipulation to make vmap work on batched gradient computation. We will fix this for all the algorithms afterward. (originally mentioned in #83 )

[TheaperDeng]

I am still not very confident about the design of this. But it's not super urgent I think. Let's wait for some feedback and make a final decision after 0.2.0