From d12b8ab882d1a629eec60e0c69f49d7f4d3b41c3 Mon Sep 17 00:00:00 2001 From: tankya2 Date: Wed, 30 Aug 2023 17:26:07 +0800 Subject: [PATCH] Black formate --- src/qibotn/cutn.py | 32 ++++++++++++++++---------------- 1 file changed, 16 insertions(+), 16 deletions(-) diff --git a/src/qibotn/cutn.py b/src/qibotn/cutn.py index 7d9984f..2fc9079 100644 --- a/src/qibotn/cutn.py +++ b/src/qibotn/cutn.py @@ -12,50 +12,50 @@ def eval(qibo_circ, datatype): def eval_tn_MPI(qibo_circ, datatype): - ncpu_threads = multiprocessing.cpu_count() // 2 n_samples = 8 - + root = 0 comm = MPI.COMM_WORLD rank = comm.Get_rank() device_id = rank % getDeviceCount() cp.cuda.Device(device_id).use() - + handle = cutn.create() cutn.distributed_reset_configuration(handle, *cutn.get_mpi_comm_pointer(comm)) network_opts = cutn.NetworkOptions(handle=handle, blocking="auto") myconvertor = QiboCircuitToEinsum(qibo_circ, dtype=datatype) operands_interleave = myconvertor.state_vector_operands() - + network = cutn.Network(*operands_interleave, options=network_opts) - network.contract_path(optimize={'samples': n_samples, 'threads': ncpu_threads}) # Calculate optimal path, returns path and info - + network.contract_path( + optimize={"samples": n_samples, "threads": ncpu_threads} + ) # Calculate optimal path, returns path and info + result = network.contract() - + cutn.destroy(handle) - + if rank == root: return result, rank if __name__ == "__main__": - - from qibo.models import QFT + from qibo.models import QFT import cupy as cp import numpy as np - - num_qubits = 10 + + num_qubits = 10 swaps = True circ_qibo = QFT(num_qubits, swaps) - - dtype="complex128" + + dtype = "complex128" sv_mpi, rank = eval_tn_MPI(circ_qibo, dtype) - + if rank == 0: sv_reference = eval(circ_qibo, dtype) state_vec = np.array(circ_qibo()) print(f"State vector difference: {abs(sv_mpi-sv_reference).max():0.3e}") assert cp.allclose(sv_mpi, sv_reference) - assert cp.allclose(sv_mpi.flatten(), state_vec) \ No newline at end of file + assert cp.allclose(sv_mpi.flatten(), state_vec)