From 34bcc5ad63e4e683793022b0e045f895ac58173e Mon Sep 17 00:00:00 2001 From: chenyu Date: Mon, 20 Jul 2026 22:42:58 -0400 Subject: [PATCH] logcumsumexp mask is bool (#17108) --- tinygrad/mixin/op.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/tinygrad/mixin/op.py b/tinygrad/mixin/op.py index 72467acac0..a5740c046b 100644 --- a/tinygrad/mixin/op.py +++ b/tinygrad/mixin/op.py @@ -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)