Explanation of the output when out of memory.

[pipme]

Hi,

I don't fully understand the BufferAssignment OOM Debugging information when an out-of-memory issue happens. After searching for a while, I didn't find any explanation in the documentation. (Especially, what is maybe_live_out allocation?)

Could you explain the following output? Thanks!

jaxlib.xla_extension.XlaRuntimeError: RESOURCE_EXHAUSTED: Out of memory while trying to allocate 14787872288 bytes.      
BufferAssignment OOM Debugging.                                                                                          
BufferAssignment stats:                                                                                                  
             parameter allocation:   27.54GiB                                                                            
              constant allocation:         0B                                                                            
        maybe_live_out allocation:   13.77GiB                                                                            
     preallocated temp allocation:         0B                                                                            
                 total allocation:   41.32GiB                                                                            
              total fragmentation:         0B (0.00%)                                                                    
Peak buffers:                 
        Buffer 1:             
                Size: 13.77GiB                              
                Entry Parameter Subshape: f64[42994,42994]                                                               
                ==========================                  
        Buffer 2:             
                Size: 13.77GiB                              
                Entry Parameter Subshape: f64[42994,42994]                                                               
                ==========================                  

        Buffer 3:             
                Size: 13.77GiB                              
                Operator: op_name="jit(fn)/jit(main)/add" source_file="/xxx/kernels.py" source_line=178       
                XLA Label: fusion                           
                Shape: f64[42994,42994]                     
                =========================

(Also, I assume my program should only require memory for an array of size [42994,42994] but it seems something is wrong. As from theBufferAssignment stats, there seem to be at least two such arrays and it tried to allocate the third one. I guess it's something I misunderstood in adding the arrays by broadcasting #18905. Have to figure it out once I understand the output stats and Peak buffers better.)