geohot/tinygrad
1
0
Fork 0
forked from tinygrad/tinygrad
Code Pull Requests Actions Activity

fix repeat_interleave with negative dim (#7734)

Browse Source
This commit is contained in:
chenyu
2024-11-16 10:15:29 -05:00
committed by GitHub
parent e3105675fb
commit f1efd84c92
2 changed files with 3 additions and 1 deletions
Download Patch File Download Diff File Expand all files Collapse all files
test/test_ops.py
+2
View File
@@ -2047,6 +2047,8 @@ class TestOps(unittest.TestCase):
helper_test_op([(3, 3)], lambda x: x.repeat_interleave(6))
helper_test_op([(3, 3)], lambda x: x.repeat_interleave(2, 1))
helper_test_op([(3, 3)], lambda x: x.repeat_interleave(2, 0))
helper_test_op([(3, 3)], lambda x: x.repeat_interleave(2, -1))
helper_test_op([(3, 3)], lambda x: x.repeat_interleave(2, -2))
def test_simple_repeat(self):
repeats = [3, 3, 4]
tinygrad/tensor.py
+1 -1
View File
@@ -1279,7 +1279,7 @@ class Tensor(SimpleMathTrait): # pylint: disable=abstract-method
print(t.repeat_interleave(2).numpy())
```
"""
x, dim = (self.flatten(), 0) if dim is None else (self, dim)
x, dim = (self.flatten(), 0) if dim is None else (self, self._resolve_dim(dim))
shp = x.shape
return x.reshape(*shp[:dim+1], 1, *shp[dim+1:]).expand(*shp[:dim+1], repeats, *shp[dim+1:]).reshape(*shp[:dim], shp[dim]*repeats, *shp[dim+1:])