Memory reusing / garbage collection mechanism during a single eval

[tyeah]

I'm curious about the memory reusing / garbage collection mechanism during a single eval. For example, for the following script:

import mlx.core as mx
import mlx.nn as nn

mx.metal.set_memory_limit(10000, relaxed=True)
mx.metal.set_cache_limit(0)
B = 100 
N = 10
D = 64

module_list = []
for _ in range(N):
    module_list.append(nn.Linear(D, D, bias=False))
    module_list.append(nn.ReLU())
model = nn.Sequential(*module_list)

x = mx.random.normal((B, D))
y = model(x)
mx.eval(y)

print(mx.metal.get_peak_memory())
print(mx.metal.get_active_memory())

I got 262200 for peak memory and 229400 for active memory. If I do the math, (262200 - 64 * 64 * 4 * 10) / (100 * 64 * 4), I get 3.8421875. So (peak memory - memory for weights) / (memory for single activation) < number of layers (10). It means some memory for intermediate activations is released / reused. But when I look at the source code, I cannot find where the destructor for array->array_desc_->data is called explicitly. And since y is still live, I also don't expect the arrays for intermediates being destructed since they are referred by y.

Could someone help me understand how this magic happens?

[tyeah]

Read deeper into the code. Is it because after calling eval on each intermediate array, arr.detach() is called, so that its inputs' ref count --, and eventually get destructed?

[awni]

That's correct!

[azhucea]

Got it, thank you!