nimlgen/tinygrad
1
0
Fork 0
forked from tinygrad/tinygrad
Code Pull Requests Activity

update linear interpolate to use int math for indices (#17441)

Browse Source
This commit is contained in:
chenyu
2026-08-06 20:18:39 -04:00
committed by GitHub
parent 1fd6b1035f
commit d8cbc11105
1 changed files with 6 additions and 4 deletions
Download Patch File Download Diff File Expand all files Collapse all files
tinygrad/mixin/op.py
+6 -4
View File
@@ -1057,14 +1057,16 @@ class OpMixin(ElementwiseMixin, ReduceMixin):
assert not (align_corners and mode != "linear"), "align_corners option can only be set with the interpolating mode linear"
x, expand = self, list(self.shape)
for i in range(-1,-len(size)-1,-1):
scale = (int(self.shape[i]) - int(align_corners)) / (size[i] - int(align_corners))
arr, reshape = type(self).arange(size[i], dtype=dtypes.float32), [1] * self.ndim
in_sz, reshape = int(self.shape[i]), [1] * self.ndim
reshape[i] = expand[i] = size[i]
if mode == "linear":
index = (scale*arr if align_corners else (scale*(arr+0.5))-0.5).clip(0, self.shape[i]-1)
low, high, perc = [y.reshape(reshape).expand(expand) for y in (index.floor().int(), index.ceil().int(), index - index.floor())]
arr = type(self).arange(size[i])
num, den = (arr*(in_sz-1), size[i]-1) if align_corners else ((arr*2+1)*in_sz - size[i], size[i]*2)
num = num.clip(0, (in_sz-1)*den)
low, high, perc = [y.reshape(reshape).expand(expand) for y in (num//den, (num+den-1)//den, (num % den).cast(dtypes.float32)/den)]
x = x.gather(i, low).lerp(x.gather(i, high), perc)
else:
scale, arr = in_sz / size[i], type(self).arange(size[i], dtype=dtypes.float32)
index = (scale*(arr+0.5) if mode=="nearest-exact" else scale*arr).cast(dtypes.int32).reshape(reshape).expand(expand)
x = x.gather(i, index)
return x.cast(self.dtype)