diff --git a/tinygrad/mixin/op.py b/tinygrad/mixin/op.py index b501faac9a..1b86cbda42 100644 --- a/tinygrad/mixin/op.py +++ b/tinygrad/mixin/op.py @@ -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)