Vmap isn't working when using JAX-based MPSCircuits

[Muzhou-Ma]

Issue Description

In the test of #213, I found that the Vmap function isn't working when using JAX-based MPSCircuit. The program is not paralleled and only uses one CPU core. The bug seems to be caused by MPSCircuit, when using JAX-based ordinary circuit, everything is fine.

[refraction-ray]

can be reproduced, may be due to the same issue for QR and SVD. These operations might not support vmap.

Updated: nope, jax can vmap qr and svd, the reason of vmap failure in MPSCircuit requires further investigation

[refraction-ray]

tf backend vmap is ok but with very low CPU utilization, only around 150% for my test example

[Muzhou-Ma]

tf backend vmap is ok but with very low CPU utilization, only around 150% for my test example

tf backend seems to have a warning with QR decomposition, perhaps this will cause low CPU utilization.