Controlling memory usage with host_callback

[davisyoshida]

Consider the following program:

import jax    
import jax.experimental.host_callback as hcb    
import jax.numpy as jnp    
import numpy as np    
    
    
shape = (1<<27,)    
def get(_):    
    return np.ones(shape, dtype=np.float32)    
    
@jax.jit    
def f():    
    result = 0    
    for _ in range(30):                                                                     
        result += jnp.sum(hcb.call(                                                         
            get,                                                                            
            None,                                                                           
            result_shape=jax.abstract_arrays.ShapedArray(shape=shape, dtype=jnp.float32)    
        ))                                                                                  
    return result                                                                           
                                                                                            
                                                                                            
def main():                                                                                 
    print(f'{f():.2e}')                                                                     
                                                                                            
if __name__ == '__main__':                                                                  
    main()

This causes an OOM on my 8GB GPU. The total size of the arrays returned from the hcb.call calls is about 16GB, but I should be able to do this computation by preventing the later callbacks from executing until the memory to hold their result is freed (in fact I expected this to be done automatically). What's the best way to do so?