[louen]

Proposed changes

This PR fixes a bug in MultiHeadAttention.createAdditiveCausalMask which occurs with float16

for dtype in [DType.float32, .bfloat16, .float16] {
    let mask = MultiHeadAttention.createAdditiveCausalMask(2, dtype: dtype)
    print(dtype, mask.asType(.float32).asArray(Float.self))
}

which will yield invalid values in float16

float32 [-0.0, -1e+09, -0.0, -0.0]
bfloat16 [-0.0, -9.9824435e+08, -0.0, -0.0]
float16 [nan, -inf, nan, nan]

due to the fact that -1e9 is not representable in float16, yielding infinities at the masked location and nan (0 * inf) in unmasked locations.

This doesn't occur in python

for dtype in [mx.float32, mx.bfloat16, mx.float16]:
    mask = nn.MultiHeadAttention.create_additive_causal_mask(2,dtype)
    print(dtype, mask.astype(mx.float32).tolist())

which correctly uses the min representable value for all types

mlx.core.float32 [[-0.0, -3.4028234663852886e+38], [-0.0, -0.0]]
mlx.core.bfloat16 [[-0.0, -3.3895313892515355e+38], [-0.0, -0.0]]
mlx.core.float16 [[-0.0, -65504.0], [-0.0, -0.0]]

.

The fix is simply to use dtype.finfo.min where available in the function.

This PR adds a unit tests that fails before the fix and succeeds after.

Checklist

Put an x in the boxes that apply.

• I have read the CONTRIBUTING document
• I have run pre-commit run --all-files to format my code / installed pre-commit prior to committing changes
• I have added tests that prove my fix is effective or that my feature works

[davidkoski]

Nice! Just waiting on CI, thanks!