tinygrad/tinygrad
2
0
Fork 3
mirror of https://github.com/tinygrad/tinygrad.git synced 2026-08-29 16:36:07 +00:00
Code Issues Packages Projects Releases Wiki Activity

fix test logcumsumexp broken devectorize=0 (#10880)

Browse Source 
* fix test logcumsumexp numerical

* lint

* Use dtypes.min instead of -1e4
This commit is contained in:
Nino Risteski
2025-06-20 20:54:50 -04:00
committed by GitHub
parent 7636d2cdc5
commit 3771cc0f77
2 changed files with 2 additions and 3 deletions
Download Patch File Download Diff File Expand all files Collapse all files
test/test_ops.py
+1 -2
View File
@@ -3,7 +3,7 @@ import numpy as np
from typing import List, Callable
import torch
import warnings
from tinygrad.helpers import getenv, IMAGE, DEBUG, CI, Context, TRANSCENDENTAL, DEVECTORIZE, OSX
from tinygrad.helpers import getenv, IMAGE, DEBUG, CI, Context, TRANSCENDENTAL, OSX
from tinygrad import Tensor, Device, dtypes
from tinygrad.tensor import _to_np_dtype
from tinygrad.device import is_dtype_supported
@@ -1546,7 +1546,6 @@ class TestOps(unittest.TestCase):
helper_test_op([()], lambda x: torch.logcumsumexp(x, dim=0), lambda x: x.logcumsumexp(), atol=1e-7, grad_atol=1e-7)
helper_test_op([()], lambda x: torch.logcumsumexp(x, dim=-1), lambda x: x.logcumsumexp(-1), atol=1e-7, grad_atol=1e-7)
@unittest.skipIf(not DEVECTORIZE, "broken without DEVECTORIZE. TODO: fix this")
def test_logcumsumexp_numerical(self):
helper_test_op(None, lambda x: torch.logcumsumexp(x, dim=0), lambda x: x.logcumsumexp(), atol=1e-7, grad_atol=1e-7, vals=[[0.0, 100.0]])
tinygrad/tensor.py
+1 -1
View File
@@ -2105,7 +2105,7 @@ class Tensor(MathTrait):
x_cummax = x_reshaped.cummax(-1).unsqueeze(-1)
x_expand = x_reshaped.unsqueeze(1).expand(*x_reshaped.shape, last_dim_size)
mask = Tensor.ones(last_dim_size, last_dim_size, requires_grad=False, device=self.device).tril().unsqueeze(0)
ret = ((x_expand - x_cummax).exp() * mask).sum(-1).log() + x_cummax.squeeze(-1)
ret = mask.where(x_expand - x_cummax, dtypes.min(self.dtype)).exp().sum(-1).log() + x_cummax.squeeze(-1)
return ret.reshape(*x.shape).transpose(-1, axis)
def argmax(self, axis=None, keepdim=False) -> Tensor: