From 3cd9a52df261edbcfd74ba8f72ca345380bb1bbd Mon Sep 17 00:00:00 2001 From: Angelos Katharopoulos Date: Tue, 21 Apr 2026 16:41:49 -0700 Subject: [PATCH] Fix ArraysCache extend (#1177) --- mlx_lm/models/cache.py | 7 +++++-- tests/test_prompt_cache.py | 8 ++++++++ 2 files changed, 13 insertions(+), 2 deletions(-) diff --git a/mlx_lm/models/cache.py b/mlx_lm/models/cache.py index d6af64d..b84c9d6 100644 --- a/mlx_lm/models/cache.py +++ b/mlx_lm/models/cache.py @@ -644,6 +644,9 @@ class ArraysCache(_BaseCache): 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 @@ class ArraysCache(_BaseCache): 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]) diff --git a/tests/test_prompt_cache.py b/tests/test_prompt_cache.py index 3a64ea3..bd1bc75 100644 --- a/tests/test_prompt_cache.py +++ b/tests/test_prompt_cache.py @@ -743,6 +743,14 @@ class TestPromptCache(unittest.TestCase): 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))