如何在MLX中实现masked_fill？使用where时返回NaN而非预期值求助

问题出在你传入mx.where的两个值类型不匹配：float('-inf')是float32类型，而0是int类型，MLX在类型转换过程中出现异常，导致mask为False的位置错误生成了nan。

解决方法 #

统一两个输入数组的数据类型，将整数0替换为浮点数0.0，或者使用形状、类型匹配的数组生成方法：

方案1：直接使用浮点类型的0

import mlx.core as mx

mask = mx.array([[False, False, True, True],
                 [False, False, True, True],
                 [False, False, True, True],
                 [False, False, True, True]], dtype=mx.bool_)

masked_tensor = mx.where(mask, mx.array(float('-inf')), mx.array(0.0))

方案2：用full_like和zeros_like保证类型/形状一致

这种方法更灵活，适合任意形状的mask：

masked_tensor = mx.where(
    mask,
    mx.full_like(mask, float('-inf'), dtype=mx.float32),
    mx.zeros_like(mask, dtype=mx.float32)
)

运行结果 #

执行后会得到你预期的输出：

array([[0., 0., -inf, -inf],
       [0., 0., -inf, -inf],
       [0., 0., -inf, -inf],
       [0., 0., -inf, -inf]], dtype=float32)

内容的提问来源于stack exchange，提问作者rtty