How can I reduce the memory usage of a min over a large intermediate array?

[nesou2]

I am not sure how to best phrase this, but I am having an issue balancing memory usage of a large, but basic, matrix calculations. For an example, suppose I have the following basic function:

# A : ndarray , shape = (10^6, 10)
def mat_product_min(b): 
   return (A @ b).min()

This will have a very large memory footprint, due to the intermediary vector of size 10^6. This isn't a concern when running with a single b, but I need to run a computation like this many thousands of times. I have a suspicion this method is not making the best use of the hardware, being memory bottle necked. We could make the memory footprint smaller by implementing this as a scan, and then finding a good balance between how much of the matrix to compute in each batch. I would imagine this is a common pattern, that is taking some associative, commutative, operator of a large vector, so perhaps others have found a good way to tackle such a situation.

My particular instance isn't exactly this, instead being a large array look-up, similar to as follows:

# A : ndarray , shape = (10^6, 10)
# B : ndarray , shape = (10^6, 10)
# C : ndarray, shape = (25, 25)
def look_up_sum_min(a_index, b_index): 
   large_intermediate_array = C[ a_index[A] , b_index[B] ]
   return large_intermediate_array .sum(axis = 1).min()

I think the same principles apply, but this example is harder to follow. One call is extremely fast (~10^-3 s), even when vmapping to fill the 40Gb of memory. If anyone has any thoughts, it would be greatly appreciated!

[f0uriest]

It sounds like implementing the incremental min using jax.lax.associative_scan might be what you're looking for?