From 3d5fea7c03a40765309cb22e005e4adb39bcea2b Mon Sep 17 00:00:00 2001 From: Connor1996 Date: Wed, 25 Mar 2026 01:52:53 -0700 Subject: [PATCH] fix(ref): widen qwen2 week2 offset type hints Signed-off-by: Connor1996 --- src/tiny_llm_ref/qwen2_week2.py | 6 +++--- 1 file changed, 3 insertions(+), 3 deletions(-) diff --git a/src/tiny_llm_ref/qwen2_week2.py b/src/tiny_llm_ref/qwen2_week2.py index 0c38ecd..5f9047f 100644 --- a/src/tiny_llm_ref/qwen2_week2.py +++ b/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)