diff --git a/test/null/test_schedule.py b/test/null/test_schedule.py index c977a31ec2..27f12864a5 100644 --- a/test/null/test_schedule.py +++ b/test/null/test_schedule.py @@ -1472,6 +1472,18 @@ class TestSchedule(unittest.TestCase): x.softmax().sum().backward() run_linear(*check_schedule(x.grad, 4)) + def test_logsumexp_backward(self): + Tensor.manual_seed(0) + x = Tensor.randn(4, 12, 64, 64).realize() + x.logsumexp(-1).sum().backward() + run_linear(*check_schedule(x.grad, 3)) + + def test_logcumsumexp_backward(self): + Tensor.manual_seed(0) + x = Tensor.randn(4, 512).realize() + x.logcumsumexp(-1).sum().backward() + run_linear(*check_schedule(x.grad, 3)) + def test_scaled_dot_product_attention_fusion(self): x, y, z, m = (Tensor.empty(32, 8, 16, 16) for _ in range(4)) out = Tensor.scaled_dot_product_attention(x, y, z, attn_mask=m) diff --git a/tinygrad/mixin/op.py b/tinygrad/mixin/op.py index e895e11bac..ff516c2578 100644 --- a/tinygrad/mixin/op.py +++ b/tinygrad/mixin/op.py @@ -658,7 +658,7 @@ class OpMixin(ElementwiseMixin, ReduceMixin): print(t.logsumexp(axis=1).numpy()) ``` """ - m = self.max(axis=axis, keepdim=True) + m = self.max(axis=axis, keepdim=True).detach() return (self - m).exp().sum(axis=axis, keepdim=keepdim).log() + (m if keepdim else m.squeeze(axis)) def _softmax(self, axis, dtype:DTypeLike|None=None) -> tuple[Self, Self, Self]: @@ -841,7 +841,7 @@ class OpMixin(ElementwiseMixin, ReduceMixin): x = self.transpose(axis, -1) 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) + x_cummax = x.cummax(-1)[0].detach() mask = type(self).ones(last_dim_size, last_dim_size, buffer=False).tril() ret = mask.where(x_unsqueezed - x_cummax.unsqueeze(-1), self.dtype.min).exp().sum(-1).log() + x_cummax return ret.transpose(-1, axis)