Repository: aesara-devs/aesara

Issue:
Memory leak in `DimShuffle`
## Description of your problem or feature request

**Please provide a minimal, self-contained, and reproducible example.**
```python
import os
import psutil
import aesara
from aesara import tensor as at


process = psutil.Process(os.getpid())
old_rss = None
n = 100_000

x = at.sharedvar.TensorSharedVariable(
    "test", at.type.dvector, np.ones(n), strict=True
)
y = x.dimshuffle([0, "x"])
f = aesara.function([], y)
for i in range(100):
    x.set_value(np.ones(n))
    f()
    if i % 10 == 0:
        new_rss = process.memory_info().rss
        if old_rss is not None:
            print(f"Python process increased RSS by {round((new_rss - old_rss) / 1024**2, 2)} MB")
        old_rss = new_rss
```

**Please provide the full traceback of any errors.**
```python
Python process increased RSS by 6.62 MB
Python process increased RSS by 7.74 MB
Python process increased RSS by 7.64 MB
Python process increased RSS by 7.68 MB
Python process increased RSS by 7.64 MB
Python process increased RSS by 7.68 MB
Python process increased RSS by 7.41 MB
Python process increased RSS by 7.63 MB
Python process increased RSS by 7.85 MB
```

**Please provide any additional information below.**

This happens when the reshape operation is compiled to C. When I run the same code setting `aesara.config.cxx = ""`, the Python process does not increase RSS.

## Versions and main components

* Aesara version: 2.4.0
* Aesara config (`python -c "import aesara; print(aesara.config)"`)
* Python version: 3.9
* Operating system: Ubuntu 18.04
* How did you install Aesara: (conda/pip) pip

Comments:

Patch:
diff --git a/aesara/tensor/c_code/dimshuffle.c b/aesara/tensor/c_code/dimshuffle.c
index d97df5254..6c67bd1bf 100644
--- a/aesara/tensor/c_code/dimshuffle.c
+++ b/aesara/tensor/c_code/dimshuffle.c
@@ -31,6 +31,8 @@ int APPLY_SPECIFIC(cpu_dimshuffle)(PyArrayObject *input, PyArrayObject **res,
   PyArrayObject *transposed_input =
       (PyArrayObject *)PyArray_Transpose(_input, &permute);
 
+  Py_DECREF(_input);
+
   PyDimMem_FREE(permute.ptr);
 
   npy_intp *res_shape = PyArray_DIMS(transposed_input);
@@ -68,7 +70,7 @@ int APPLY_SPECIFIC(cpu_dimshuffle)(PyArrayObject *input, PyArrayObject **res,
   *res = (PyArrayObject *)PyArray_Newshape(transposed_input, &reshape_shape,
                                            NPY_CORDER);
 
-  /* Py_XDECREF(transposed_input); */
+  Py_DECREF(transposed_input);
 
   PyDimMem_FREE(reshape_shape.ptr);
 


Test patch:
diff --git a/tests/tensor/test_elemwise.py b/tests/tensor/test_elemwise.py
index 0ad4a76a6..1ad8ca9b2 100644
--- a/tests/tensor/test_elemwise.py
+++ b/tests/tensor/test_elemwise.py
@@ -1,4 +1,5 @@
 import math
+import tracemalloc
 from copy import copy
 
 import numpy as np
@@ -141,6 +142,41 @@ class TestDimShuffle(unittest_tools.InferShapeTester):
             # Confirm the broadcasted value in the output
             assert np.array_equiv(outputs[0].storage[0], 2039)
 
+    @pytest.mark.parametrize("inplace", [True, False])
+    def test_memory_leak(self, inplace):
+        import gc
+
+        n = 100_000
+
+        x = aesara.shared(np.ones(n, dtype=np.float64))
+
+        y = x.dimshuffle([0, "x"])
+        y.owner.op.inplace = inplace
+
+        f = aesara.function([], y, mode=Mode(optimizer=None))
+
+        assert len(f.maker.fgraph.apply_nodes) == 2
+        assert isinstance(f.maker.fgraph.toposort()[0].op, DimShuffle)
+
+        assert f.maker.fgraph.toposort()[0].op.inplace is inplace
+
+        tracemalloc.start()
+
+        blocks_last = None
+        block_diffs = []
+        for i in range(50):
+            x.set_value(np.ones(n))
+            _ = f()
+            _ = gc.collect()
+            blocks_i, _ = tracemalloc.get_traced_memory()
+            if blocks_last is not None:
+                blocks_diff = (blocks_i - blocks_last) // 10 ** 3
+                block_diffs.append(blocks_diff)
+            blocks_last = blocks_i
+
+        tracemalloc.stop()
+        assert np.allclose(np.mean(block_diffs), 0)
+
 
 class TestBroadcast:
     # this is to allow other types to reuse this class to test their ops