Understanding the memory usage during broadcasting

[pipme]

In the following example, the jaxprs for f1 and f2 look the same but their peak memory usages are different. Could you explain what happened?

Furthermore, is it possible to reduce the memory usage to ~4008192 bytes which corresponds to a float32 array of size (1000, 1000), and in the same time keep the function gradable and efficient?

Thanks!

import jax.numpy as jnp
import jax


def f1(X, Y):
    dist = X @ Y
    dist += X + Y
    return dist


def f2(X, Y):
    dist = X @ Y + X + Y
    return dist


X = jnp.ones((1000, 1))
Y = jnp.ones((1, 1000))
print(jax.make_jaxpr(f1)(X, Y))
print(jax.make_jaxpr(f2)(X, Y))


fun = "f1"
if fun == "f1":
    dist = f1(X, Y)
    print(jax.devices()[0].memory_stats())
    # {'bytes_in_use': 4008192, 'bytes_limit': 63768084480, 'bytes_reserved': 0, 'largest_alloc_size': 4000000, 'largest_free_block_bytes': 0, 'num_allocs': 9,
    # 'peak_bytes_in_use': 12008192, 'peak_bytes_reserved': 0, 'peak_pool_bytes': 63768084480, 'pool_bytes': 63768084480}
else:
    dist = f2(X, Y)
    print(jax.devices()[0].memory_stats())
    # {'bytes_in_use': 4008192, 'bytes_limit': 63768084480, 'bytes_reserved': 0, 'largest_alloc_size': 4000000, 'largest_free_block_bytes': 0, 'num_allocs': 9, 
    # 'peak_bytes_in_use': 8008192, 'peak_bytes_reserved': 0, 'peak_pool_bytes': 63768084480, 'pool_bytes': 63768084480}

[jakevdp]

The two versions are subtley different in their order of operations, and this leads to different temporary variables. It might help to write out the Python intermediate variables explicitly:

def f1(X, Y):
    dist = X @ Y
    tmp = X + Y
    dist = dist + tmp
    return dist

def f2(X, Y):
    dist = X @ Y
    dist = dist + X
    dist = dist + Y
    return dist

Note that while X and Y are size 1000, both X @ Y and X + Y have shape [1000, 1000]. In the first version, you create two [1000, 1000] arrays (dist and tmp), while in the second version you create only one such array (dist), and so the first version will have roughly twice the memory footprint of the second.

Does that make sense?

[pipme]

Hi @jakevdp , thanks for the reply! It's very clear. However, if you look at peak_bytes_in_use, the second version uses 1.5 times more memory than the first version. The peak_bytes_in_use in the first version corresponds to creating three float32 [1000, 1000] arrays while the second version corresponds to creating two arrays.

In other words, it seems in the second version there are still two intermediate arrays created but ideally it should only create one. Is it possible to create only one array of size [1000, 1000] without another intermediate array such that the peak_bytes_in_use will be reduced to ~4008192?