Make rank class attribute

This commit is contained in:
tankya2
2025-08-26 17:11:08 +08:00
parent 39dad0fd88
commit 65a04c32fa

View File

@@ -22,6 +22,7 @@ class CuTensorNet(QibotnBackend, NumpyBackend): # pragma: no cover
self.configure_tn_simulation(runcard) self.configure_tn_simulation(runcard)
def configure_tn_simulation(self, runcard): def configure_tn_simulation(self, runcard):
self.rank = None
if runcard is not None: if runcard is not None:
self.MPI_enabled = runcard.get("MPI_enabled", False) self.MPI_enabled = runcard.get("MPI_enabled", False)
self.NCCL_enabled = runcard.get("NCCL_enabled", False) self.NCCL_enabled = runcard.get("NCCL_enabled", False)
@@ -106,8 +107,8 @@ class CuTensorNet(QibotnBackend, NumpyBackend): # pragma: no cover
and self.NCCL_enabled == False and self.NCCL_enabled == False
and self.expectation_enabled == False and self.expectation_enabled == False
): ):
state, rank = eval.dense_vector_tn_MPI(circuit, self.dtype, 32) state, self.rank = eval.dense_vector_tn_MPI(circuit, self.dtype, 32)
if rank > 0: if self.rank > 0:
state = np.array(0) state = np.array(0)
elif ( elif (
self.MPI_enabled == False self.MPI_enabled == False
@@ -115,8 +116,8 @@ class CuTensorNet(QibotnBackend, NumpyBackend): # pragma: no cover
and self.NCCL_enabled == True and self.NCCL_enabled == True
and self.expectation_enabled == False and self.expectation_enabled == False
): ):
state, rank = eval.dense_vector_tn_nccl(circuit, self.dtype, 32) state, self.rank = eval.dense_vector_tn_nccl(circuit, self.dtype, 32)
if rank > 0: if self.rank > 0:
state = np.array(0) state = np.array(0)
elif ( elif (
self.MPI_enabled == False self.MPI_enabled == False
@@ -131,10 +132,10 @@ class CuTensorNet(QibotnBackend, NumpyBackend): # pragma: no cover
and self.NCCL_enabled == False and self.NCCL_enabled == False
and self.expectation_enabled == True and self.expectation_enabled == True
): ):
state, rank = eval.expectation_tn_MPI( state, self.rank = eval.expectation_tn_MPI(
circuit, self.dtype, self.observable, 32 circuit, self.dtype, self.observable, 32
) )
if rank > 0: if self.rank > 0:
state = np.array(0) state = np.array(0)
elif ( elif (
self.MPI_enabled == False self.MPI_enabled == False
@@ -142,10 +143,10 @@ class CuTensorNet(QibotnBackend, NumpyBackend): # pragma: no cover
and self.NCCL_enabled == True and self.NCCL_enabled == True
and self.expectation_enabled == True and self.expectation_enabled == True
): ):
state, rank = eval.expectation_tn_nccl( state, self.rank = eval.expectation_tn_nccl(
circuit, self.dtype, self.observable, 32 circuit, self.dtype, self.observable, 32
) )
if rank > 0: if self.rank > 0:
state = np.array(0) state = np.array(0)
else: else:
raise_error(NotImplementedError, "Compute type not supported.") raise_error(NotImplementedError, "Compute type not supported.")