test: fully mask idle batching slots (#109)

Signed-off-by: Connor1996 <zbk602423539@gmail.com>

Author: Connor1996

src/tiny_llm_ref/kv_cache.py

@@ -78,9 +78,6 @@ def get_seq_len(data):
         )
         for b in range(B):
             if data[b] is None:
-                # for some reasons we need to do this, otherwise it will cause wrong output?
-                # maybe precision issues?
-                masks[b, :, :] = causal_mask(mask_length, seq_len, dtype=key.dtype)
                 continue
             key, value, S, mask = data[b]
             keys[b, :, seq_len - S : seq_len, :] = key

tests_refsol/test_week_2_day_6.py

@@ -236,7 +236,7 @@ def test_task_2_batching_kv_cache():
     expected_mask = mx.array(
         [
             [[[-mx.inf, 0.0, 0.0, -mx.inf], [-mx.inf, 0.0, 0.0, 0.0]]],
-            [[[0.0, 0.0, 0.0, -mx.inf], [0.0, 0.0, 0.0, 0.0]]],
+            [[[-mx.inf, -mx.inf, -mx.inf, -mx.inf], [-mx.inf, -mx.inf, -mx.inf, -mx.inf]]],
             [[[0.0, 0.0, 0.0, -mx.inf], [0.0, 0.0, 0.0, 0.0]]],
         ],
         dtype=mx.float32,