torch:masked-fill

(torch:masked-fill a mask value)

Differentiable masked fill (torch.masked_fill over linalg:where): the scalar value where mask is non-zero, a's element where it is zero. mask (a 0/1 array, a comparison mask, or a tensor) and value are constants -- no gradient flows to them; a's gradient is zero at the filled positions. Filling attention scores with -infinity before torch:softmax is the masked-attention idiom.