zhangyiss/mlx
1
0
Fork 0
mirror of https://github.com/ml-explore/mlx.git synced 2025-09-18 18:28:12 +08:00
Code Issues Actions Packages Projects Releases Wiki Activity

Add mx.finfo and use it when making causal mask (#1726)

Browse Source 
* finfo

* fixes

* docs
This commit is contained in:
Awni Hannun
2024-12-19 14:52:41 -08:00
committed by GitHub
parent e03f0372b1
commit c3628eea49
9 changed files with 154 additions and 3 deletions
Download Patch File Download Diff File Expand all files Collapse all files

4
python/mlx/nn/layers/transformer.py
View File

@@ -102,9 +102,7 @@ class MultiHeadAttention(Module):
def create_additive_causal_mask(N: int, dtype: mx.Dtype = mx.float32):
indices = mx.arange(N)
mask = indices[:, None] < indices[None]
# usually inf but 1e9 is as good and softmax(full(1e9)) != nan
# TODO: Should replace this with finfo(dtype).min
mask = mask.astype(dtype) * -1e9
mask = mask.astype(dtype) * mx.finfo(dtype).min
return mask