<SYSTEM_TASK:>
Given the following text description, write Python code to implement the functionality described below step by step
<END_TASK>
<USER_TASK:>
Description:
 The RadixSort.argsort() sorts the input array and outputs an array containing the indices of the sort. In the end the input array is sorted!
Step1:  sl mst lifetime gpu
Step2:  Compute the MST for the input graph.
Step3:  Allocate array for the mst weights.
Step4:  Get array with only the considered weights in the MST.
Step5:  Sort the MST weights. There are no repeted edges at this point since the output MST is like a directed graph.
Step6:  Allocate array for the lifetimes.
Step8:  Between the above cell and the next, I have to get the argmax of lifetimes. I also have to sort the mst with the argsort from the weights. Don't forget to update the number of n_edges. Send those arrays to build the resulting graph.
<ASSISTANT_TASK:>
Python Code:
4e7*4/1024/1024
a=np.random.randint(0,1e4,1e6)
dA = cuda.to_device(a)
del a
@cuda.reduce
def argmax_gpu(a,b):
    if a >= b:
        return a
    else:
        return b
%time a.max()
%time argmax_gpu(dA)
sorter = RadixSort(maxcount=dA.size, dtype=dA.dtype)
dRes = sorter.argsort(dA)
res = dRes.copy_to_host()
a
res
Explanation: The RadixSort.argsort() sorts the input array and outputs an array containing the indices of the sort. In the end the input array is sorted!
End of explanation
weight = dA
MAX_TPB = 256
myStream = 0
Explanation: sl mst lifetime gpu
End of explanation
# get MST
mst, n_edges = boruvka_minho_gpu(dest, weight, firstEdge, outDegree, MAX_TPB=MAX_TPB)
Explanation: Compute the MST for the input graph.
End of explanation
h_n_edges = n_edges.getitem(0)
mst_weights = cuda.device_array(h_n_edges, dtype = weight.dtype)
Explanation: Allocate array for the mst weights.
End of explanation
mstGrid = compute_cuda_grid_dim(h_n_edges, MAX_TPB)
getWeightsOfEdges_gpu[mstGrid, MAX_TPB, myStream](mst, n_edges, weight, mst_weights)
Explanation: Get array with only the considered weights in the MST.
End of explanation
RadixSort.sort
sorter = RadixSort(maxcount = mst_weights.size, dtype = mst_weights.dtype)
sortedWeightArgs = sorter.argsort(mst_weights)
Explanation: Sort the MST weights. There are no repeted edges at this point since the output MST is like a directed graph.
End of explanation
lifetimes = cuda.device_array(mst_weights.size - 1, dtype = nweight.dtype)
compute_lifetimes_CUDA[lifetimeGrid, MAX_TPB, myStream](nweight, lifetimes)
Explanation: Allocate array for the lifetimes.
End of explanation
# get contracted graph from MST
ndest, nweight, nfe, nod = getGraphFromEdges_gpu(dest, weight, fe, od, edges = mst, n_edges = n_edges,\
                                                 MAX_TPB = MAX_TPB, stream = myStream)
nweight = dA
lifetimeGrid = compute_cuda_grid_dim(lifetimes.size, MAX_TPB)
@cuda.jit
def getWeightsOfEdges_gpu(edges, n_edges, weights, nweights):
    
    This function will take a list of edges (edges), the number of edges to consider (n_edges,
    the weights of all the possible edges (weights) and the array for the weights of the list
    of edges and put the weight of each edge in the list of edges in the nweights, in the same
    position.
    
    n_edges_sm = cuda.shared.array(1, dtype = np.int32)
    edge = cuda.grid(1)
    
    if edge == 0:
        n_edges_sm[0] = n_edges[0]
    
    if edge >= edges.size:
        return
    
    cuda.syncthreads()
    
    nweights[edge] = weights[edges[edge]]
@cuda.jit
def compute_lifetimes_CUDA(nweight, lifetimes):
    edge = cuda.grid(1)
    
    if edge >= lifetimes.size:
        return
    
    lifetimes[edge] = nweight[edge + 1] - nweight[edge]
res = sortedWeightArgs.copy_to_host()
res
Explanation: Between the above cell and the next, I have to get the argmax of lifetimes. I also have to sort the mst with the argsort from the weights. Don't forget to update the number of n_edges. Send those arrays to build the resulting graph.
End of explanation
<END_TASK>