fix multi output leak (#1548)

Author: awni

mlx/array.cpp

@@ -271,6 +271,9 @@ array::ArrayDesc::~ArrayDesc() {
     for (array& a : ad.inputs) {
       if (a.array_desc_) {
         input_map.insert({a.id(), a});
+        for (auto& s : a.siblings()) {
+          input_map.insert({s.id(), s});
+        }
       }
     }
     ad.inputs.clear();

python/tests/test_array.py

@@ -1771,6 +1771,19 @@ def fun():
         peak_2 = mx.metal.get_peak_memory()
         self.assertEqual(peak_1, peak_2)
 
+        def fun():
+            a = mx.array([1.0, 2.0, 3.0, 4.0])
+            b, _ = mx.divmod(a, a)
+            return mx.log(b)
+
+        fun()
+        mx.synchronize()
+        peak_1 = mx.metal.get_peak_memory()
+        fun()
+        mx.synchronize()
+        peak_2 = mx.metal.get_peak_memory()
+        self.assertEqual(peak_1, peak_2)
+
     def test_add_numpy(self):
         x = mx.array(1)
         y = np.array(2, dtype=np.int32)