From c15c926b75914700ce8fcf68ee89e4cdd4245420 Mon Sep 17 00:00:00 2001 From: dikanquit Date: Fri, 29 May 2026 11:21:29 +0200 Subject: [PATCH] diagonal: clamp out-of-range offset to return empty |offset| > dim raised ValueError from shrink --- tinygrad/mixin/movement.py | 1 + 1 file changed, 1 insertion(+) diff --git a/tinygrad/mixin/movement.py b/tinygrad/mixin/movement.py index dc03177ec9d4b..575f7b5ee83c1 100644 --- a/tinygrad/mixin/movement.py +++ b/tinygrad/mixin/movement.py @@ -500,6 +500,7 @@ def diagonal(self, offset:int=0, dim1:int=0, dim2:int=1) -> Self: """ if (dim1:=self._resolve_dim(dim1)) == (dim2:=self._resolve_dim(dim2)): raise RuntimeError("dim1 and dim2 cannot be the same dimension") x = self.permute(*[i for i in range(self.ndim) if i != dim1 and i != dim2], dim1, dim2) + offset = max(-int(x.shape[-2]), min(int(x.shape[-1]), offset)) if offset >= 0: x = x.shrink(tuple(None for _ in x.shape[:-1]) + ((offset, x.shape[-1]),)) else: x = x.shrink(tuple(None for _ in x.shape[:-2]) + ((-offset, x.shape[-2]), None)) if (d := min(int(x.shape[-2]), int(x.shape[-1]))) <= 0: return x.reshape(*x.shape[:-2], 0)