Fix ArraysCache extend (#1177)

Author: angeloskath

mlx_lm/models/cache.py

@@ -644,6 +644,9 @@ def extend(self, other):
         In-place extend this cache with the other cache.
         """
 
+        a_batch = self.batch_size
+        b_batch = other.batch_size
+
         def cat(a, b):
             shape = dtype = None
             if a is not None:
@@ -657,9 +660,9 @@ def cat(a, b):
                 return None
 
             if a is None:
-                a = mx.zeros((self.batch_size,) + shape[1:], dtype=dtype)
+                a = mx.zeros((a_batch,) + shape[1:], dtype=dtype)
             if b is None:
-                b = mx.zeros((other.batch_size,) + shape[1:], dtype=dtype)
+                b = mx.zeros((b_batch,) + shape[1:], dtype=dtype)
 
             return mx.concatenate([a, b])
 

tests/test_prompt_cache.py

@@ -743,6 +743,14 @@ def test_arrays_cache_extend_with_empty(self):
         self.assertEqual(empty2[0].shape, (4, 4, 8))
         self.assertEqual(empty2[1].shape, (4, 4))
 
+        # Extend content with empty
+        content = ArraysCache.merge((c1, c2))
+        empty2 = ArraysCache.merge((ArraysCache(2), ArraysCache(2)))
+        content.extend(empty2)
+        self.assertEqual(content[0].shape, (4, 4, 8))
+        self.assertEqual(content[1].shape, (4, 4))
+        self.assertEqual(content.make_mask(10).shape, (4, 10))
+
         # multiple empty extensions accumulate correctly
         stepwise = ArraysCache.merge((c1,))
         stepwise.extend(ArraysCache(2))