From 6228130be46fe2a6d9e9b818b41cad84195c25c5 Mon Sep 17 00:00:00 2001 From: Chen-Yu Yang Date: Fri, 15 Nov 2024 23:11:04 -0500 Subject: [PATCH] remove add in pad when padding with non-zero this is simpler before shapetracker is shared --- test/external/external_test_opt.py | 2 +- tinygrad/tensor.py | 2 +- 2 files changed, 2 insertions(+), 2 deletions(-) diff --git a/test/external/external_test_opt.py b/test/external/external_test_opt.py index 584c3af89f4bc..b20bf8bf279a7 100644 --- a/test/external/external_test_opt.py +++ b/test/external/external_test_opt.py @@ -68,7 +68,7 @@ def test_resnet(self): model = ResNet18() for p in get_parameters(model): p.assign(np.zeros(p.shape, dtype=_to_np_dtype(p.dtype))) img = Tensor.randn(1, 3, 224, 224) - with CLCache(23): + with CLCache(24): model.forward(img).realize() def test_vit(self): diff --git a/tinygrad/tensor.py b/tinygrad/tensor.py index 0efb61ad9b74e..5c56e0a97ba99 100644 --- a/tinygrad/tensor.py +++ b/tinygrad/tensor.py @@ -1039,7 +1039,7 @@ def pad(self, padding:Union[Sequence[sint], Sequence[Optional[Tuple[sint, sint]] pX = ((0,0),)*(self.ndim - len(padding)//2) + tuple(zip(padding[-2::-2], padding[::-2])) if flat else padding if len(pX) != self.ndim: raise ValueError(f"padding length is improper, {padding=} {self.ndim=}") X, pX = self, cast(Tuple[Tuple[sint, sint]], tuple((0,0) if p is None else p for p in pX)) - def _constant(x,px,v): return F.Pad.apply(x, arg=px) if v == 0 else F.Pad.apply(x, arg=px) + F.Pad.apply(Tensor.ones_like(x), arg=px).where(0, v) + def _constant(x,px,v): return F.Pad.apply(x, arg=px) if v == 0 else F.Pad.apply(Tensor.ones_like(x), arg=px).where(F.Pad.apply(x, arg=px), v) # early return for symbolic with positive pads (no need to max) if all(resolve(p >= 0) for p in flatten(pX)): return _constant(X, pX, value) pads, shrinks = tuple((smax(pB,0), smax(pA,0)) for pB,pA in pX), tuple((-smin(pB,0),smin(pA+s,s)) for (pB,pA),s in zip(pX, self.shape))