From a54dfd698e70dc3ef2b79b80b05968321abc1a05 Mon Sep 17 00:00:00 2001 From: Yi Wang Date: Mon, 10 Jun 2024 15:18:34 -0700 Subject: [PATCH] Correct the type annotation of cache in llama.py (#828) * Update * Fix isort --- llms/mlx_lm/models/llama.py | 6 +++--- 1 file changed, 3 insertions(+), 3 deletions(-) diff --git a/llms/mlx_lm/models/llama.py b/llms/mlx_lm/models/llama.py index e7f4f16..2a49ee3 100644 --- a/llms/mlx_lm/models/llama.py +++ b/llms/mlx_lm/models/llama.py @@ -4,7 +4,7 @@ from typing import Dict, Optional, Tuple, Union import mlx.core as mx import mlx.nn as nn -from .base import BaseModelArgs, create_additive_causal_mask +from .base import BaseModelArgs, KVCache, create_additive_causal_mask @dataclass @@ -73,7 +73,7 @@ class Attention(nn.Module): self, x: mx.array, mask: Optional[mx.array] = None, - cache: Optional[Tuple[mx.array, mx.array]] = None, + cache: Optional[KVCache] = None, ) -> mx.array: B, L, D = x.shape @@ -135,7 +135,7 @@ class TransformerBlock(nn.Module): self, x: mx.array, mask: Optional[mx.array] = None, - cache: Optional[Tuple[mx.array, mx.array]] = None, + cache: Optional[KVCache] = None, ) -> mx.array: r = self.self_attn(self.input_layernorm(x), mask, cache) h = x + r