400 lines

test/unit/test_attention.py

@@ -16,10 +16,10 @@ def apply_rope(x:Tensor, start_pos:int):
 class TestLinear(unittest.TestCase):
   def test_recovers_packed_ggml_weight(self):
     for ggml_type,packed_size in ((13, 176), (14, 210), (23, 136)):
-      packed = Tensor.empty(packed_size+1, dtype=dtypes.uint8, device="CPU")[1:]
+      packed = Tensor.empty(packed_size+4, dtype=dtypes.uint8, device="CPU")[4:]
       decoded = ggml_data_to_tensor(packed, 256, ggml_type).reshape(1, 256)
       linear = Linear(256, 1, bias=False)
-      self.assertTrue(linear.set_quantized(decoded))
+      self.assertIsNotNone(linear.set_quantized(decoded))
       self.assertEqual((linear.ggml_type, linear.weight.shape), (ggml_type, (packed_size,)))
 
 class TestAttention(unittest.TestCase):
@@ -77,9 +77,8 @@ class TestGatedDeltaNetBlock(unittest.TestCase):
   def _run_attention(self, block:GatedDeltaNetBlock, x:Tensor, start_pos:int):
     x_norm = block.attn_norm(x)
     block._init_state(x_norm)
-    out = block._attention(x_norm, start_pos).realize()
-    if block.pending_state is not None:
-      Tensor.realize(block.conv_state.assign(block.pending_state[0]), block.recurrent_state.assign(block.pending_state[1]))
+    out, conv_state, recurrent_state = block._attention(x_norm, start_pos)
+    Tensor.realize(out, block.conv_state.assign(conv_state), block.recurrent_state.assign(recurrent_state))
     return out.numpy()
 
   def _cache_views(self, block:GatedDeltaNetBlock) -> tuple[np.ndarray, np.ndarray]:

test/unit/test_function.py

@@ -67,8 +67,9 @@ class TestFunction(unittest.TestCase):
     a, b = Tensor.zeros(8).contiguous().realize(), Tensor.zeros(8).contiguous().realize()
     @function(precompile=True, allow_implicit=True)
     def f(x:Tensor, start:UOp):
-      stores = (a[start:start+2].uop.store(x.uop), b[start:start+2].uop.store((x+1).uop))
-      return Tensor(a.uop.after(*stores)) + Tensor(b.uop.after(*stores))
+      a[start:start+2].assign(x)
+      b[start:start+2].assign(x+1)
+      return a+b
     out = f(Tensor([2., 3.]).realize(), UOp.variable("start", 0, 6).bind(1))
     np.testing.assert_equal(out.numpy(), [0, 5, 7, 0, 0, 0, 0, 0])