import numpy as np
from mpi4py import MPI
import sys
import time

#Sets the MPI variables
comm = MPI.COMM_WORLD
size = comm.Get_size()
rank = comm.Get_rank()

v_size = 2**28
v_size_local = int(v_size/size)


data = np.random.rand(v_size_local)
start_time = time.time()

data_mm = np.dot(data, data)

data_gather = comm.gather(data_mm,root = 0)

end_time = time.time()
if(rank == 0):
    data_sum = np.sum(data_gather)
    print(f"{size} {data.nbytes*size} {end_time-start_time}")