fix jac with vmap

Author: refraction-ray

CHANGELOG.md

@@ -16,6 +16,8 @@
 
 - Fix customized jax eigh operator by noting the return is a namedtuple
 
+- Fix some issues in `jacfwd` and `jacrev` when integrated with vmap
+
 ## 1.0.2
 
 ### Added

tensorcircuit/backends/abstract_backend.py

@@ -1548,7 +1548,7 @@ def wrapper(*args: Any, **kws: Any) -> Any:
                                 if i == argnum
                                 else self.reshape(
                                     self.zeros(
-                                        [self.sizen(arg), self.sizen(arg)],
+                                        [self.sizen(args[argnum]), self.sizen(arg)],
                                         dtype=arg.dtype,
                                     ),
                                     [-1] + list(self.shape_tuple(arg)),
@@ -1636,6 +1636,7 @@ def _first(x: Sequence[Any]) -> Any:
                     ),
                     jj,
                 )
+                jj = [jji for ind, jji in enumerate(jj) if ind in argnums]
                 if len(jj) == 1:
                     jj = jj[0]
                 jjs.append(jj)