Is there a way to compute in parallel the maximum value of an array so big I need to use a generator in parallel? I am computing the maximum value of a function over a very large discrete set of matrices (9 billion elements) but the sequential code takes around 200 hours to run on Google colab’s CPU.

Going into more detail, I have a 680×4 array and need to compute a function taking in four rows of the array, not allowing duplicate rows and where order doesn’t matter. From these rows I create a 4×4 matrix and evaluate my function. This gives 680 choose 4 ~ 9 billion possible matrices. I am currently using a generator that returns each possible combination, creating the matrix, applying my function, and keeping track of the maximum value. I imagine there must be a nice way of parallelizing on GPU. I was also thinking about using JAX but I am not familiar with parallel computing at all. Any pointers are appreciated.

Here’s a summary of for loop code. The rows are given.

max_value = 0

for combination in tqdm(combinations_generator, total = total_iterations):

        A = rows[torch.tensor(combination)] # Create matrix with chosen rows

        value_A = function(A) # compute function of the matrix

        if value_A > max_value: # keep track of maximum value
            value_A, best_A = value_A, A