search

tencent-quantum-lab / tensorcircuit

Tensor network based quantum software framework for the NISQ era
https://tensorcircuit.readthedocs.io
Apache License 2.0
252 stars 75 forks source link

Vmap isn't working when using JAX-based MPSCircuits #214

Open Marsmmz opened 1 month ago

Marsmmz commented 1 month ago

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 commented 1 month ago

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 commented 1 month ago

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

Marsmmz commented 1 month ago

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.