zhangyiss/mlx-examples
1
0
Fork 0
mirror of https://github.com/ml-explore/mlx-examples.git synced 2025-10-23 22:18:06 +08:00
Code Issues Actions Packages Projects Releases Wiki Activity

Change tuple type definitions to use Tuple (#308)

Browse Source
This commit is contained in:
Angelos Katharopoulos
2024-01-12 11:15:09 -08:00
committed by GitHub
parent c1342b8e89
commit 1fa40067fe
3 changed files with 6 additions and 5 deletions
Download Patch File Download Diff File Expand all files Collapse all files

3
llms/mlx_lm/models/phi2.py
View File

@@ -1,5 +1,6 @@
import math
from dataclasses import dataclass
from typing import Tuple
import mlx.core as mx
import mlx.nn as nn
@@ -128,7 +129,7 @@ class Model(nn.Module):
x: mx.array,
mask: mx.array = None,
cache: mx.array = None,
) -> tuple[mx.array, mx.array]:
) -> Tuple[mx.array, mx.array]:
mask = None
if x.shape[1] > 1:
mask = nn.MultiHeadAttention.create_additive_causal_mask(x.shape[1])