fix(ref): widen qwen2 week2 offset type hints

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

Author: Connor1996

src/tiny_llm_ref/qwen2_week2.py

@@ -53,7 +53,7 @@ def __init__(
     def __call__(
         self,
         x: mx.array,
-        offsets: list[int],
+        offsets: int | list[int] | mx.array,
         cache: TinyKvCache,
         mask: mx.array | str | None = None,
     ) -> mx.array:
@@ -172,7 +172,7 @@ def __init__(
     def __call__(
         self,
         x: mx.array,
-        offset: int,
+        offset: int | list[int] | mx.array,
         cache: TinyKvCache,
         mask: mx.array | str | None = None,
     ) -> mx.array:
@@ -266,7 +266,7 @@ def __init__(
     def __call__(
         self,
         inputs: mx.array,
-        offset: int,
+        offset: int | list[int] | mx.array,
         cache: list[TinyKvCache],
     ) -> mx.array:
         h = self.embedding(inputs)