logcumsumexp mask is bool (#17108)

2026-07-20 22:42:58 -04:00
parent 40f0d4af14
commit 34bcc5ad63
1 changed files with 1 additions and 1 deletions
@@ -839,7 +839,7 @@ class OpMixin(ElementwiseMixin, ReduceMixin):
    last_dim_size = x.shape[-1]
    x_unsqueezed = x.unsqueeze(-2).expand((None,)*(self.ndim-1)+(last_dim_size, None))
    x_cummax = x.cummax(-1)[0].detach()
-    mask = type(self).ones(last_dim_size, last_dim_size, buffer=False).tril()
+    mask = type(self).ones(last_dim_size, last_dim_size, buffer=False, dtype=dtypes.bool).tril()
    ret = mask.where(x_unsqueezed - x_cummax.unsqueeze(-1), self.dtype.min).exp().sum(-1).log() + x_cummax
    return ret.transpose(-1, axis)