add test for week 1 day 5 test 1: Qwen2TransformerBlock (#59)

Author: jiengup

tests_refsol/test_week_1_day_5.py

@@ -3,8 +3,81 @@
 from .tiny_llm_base import Qwen2ModelWeek1, Embedding, dequantize_linear, qwen2_week1
 from mlx_lm import load
 
-# TODO: task 1 tests
-
+@pytest.mark.parametrize("stream", AVAILABLE_STREAMS, ids=AVAILABLE_STREAMS_IDS)
+@pytest.mark.parametrize("precision", PRECISIONS, ids=PRECISION_IDS)
+@pytest.mark.parametrize("mask", [None, "causal"], ids=["no_mask", "causal_mask"])
+def test_task_1_transformer_block(
+    stream: mx.Stream, precision: mx.Dtype, mask: str | None
+):
+    with mx.stream(stream):
+        from mlx_lm.models import qwen2
+
+        BATCH_SIZE = 1
+        SEQ_LEN = 10
+        NUM_ATTENTION_HEAD = 4
+        NUM_KV_HEADS = 2
+        HIDDEN_SIZE = 32
+        INTERMEDIATE_SIZE = HIDDEN_SIZE * 4
+        
+        args = qwen2.ModelArgs(
+            model_type="qwen2",
+            hidden_size=HIDDEN_SIZE,
+            num_hidden_layers=1,
+            intermediate_size=INTERMEDIATE_SIZE,
+            num_attention_heads=NUM_ATTENTION_HEAD,
+            num_key_value_heads=NUM_KV_HEADS,
+            rms_norm_eps=1e-6,
+            vocab_size=1000,
+        )
+
+        mlx_transformer_block = qwen2.TransformerBlock(args)
+
+        mlx_attention = mlx_transformer_block.self_attn
+        wq = mlx_attention.q_proj.weight
+        wk = mlx_attention.k_proj.weight
+        wv = mlx_attention.v_proj.weight
+        wo = mlx_attention.o_proj.weight
+        bq = mlx_attention.q_proj.bias
+        bk = mlx_attention.k_proj.bias
+        bv = mlx_attention.v_proj.bias
+
+        mlx_mlp = mlx_transformer_block.mlp
+        w_gate = mlx_mlp.gate_proj.weight
+        w_up = mlx_mlp.up_proj.weight
+        w_down = mlx_mlp.down_proj.weight
+
+        w_input_layernorm = mlx_transformer_block.input_layernorm.weight
+        w_post_attention_layernorm = mlx_transformer_block.post_attention_layernorm.weight
+        
+        user_transformer_block = qwen2_week1.Qwen2TransformerBlock(
+            num_attention_heads=NUM_ATTENTION_HEAD,
+            num_kv_heads=NUM_KV_HEADS,
+            hidden_size=HIDDEN_SIZE,
+            intermediate_size=INTERMEDIATE_SIZE,
+            rms_norm_eps=1e-6,
+            wq=wq,
+            wk=wk,
+            wv=wv,
+            wo=wo,
+            bq=bq,
+            bk=bk,
+            bv=bv,
+            w_gate=w_gate,
+            w_up=w_up,
+            w_down=w_down,
+            w_input_layernorm=w_input_layernorm,
+            w_post_attention_layernorm=w_post_attention_layernorm
+        )
+
+        mx.random.seed(42)
+        x = mx.random.uniform(
+            shape=(BATCH_SIZE, SEQ_LEN, HIDDEN_SIZE), dtype=precision
+        )
+
+        user_output = user_transformer_block(x, mask=mask)
+        mlx_output = mlx_transformer_block(x, mask=mask, cache=None)
+        
+        assert_allclose(user_output, mlx_output, precision=precision, rtol=1e-1)
 
 @pytest.mark.skipif(
     not qwen_2_05b_model_exists(), reason="Qwen2-0.5B-Instruct-MLX model not found"