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

check arg to Tensor.flip can appear only once (#5068)

Browse Source 
* check arg to Tensor.flip can appear only once

raise RuntimeError if there are multiple

* fix test
This commit is contained in:
chenyu
2024-06-20 09:33:42 -04:00
committed by GitHub
parent 55e02cdd84
commit e8f39fcaaa
3 changed files with 12 additions and 7 deletions
Download Patch File Download Diff File Expand all files Collapse all files
test/test_ops.py
+7 -4
View File
@@ -1122,9 +1122,9 @@ class TestOps(unittest.TestCase):
helper_test_op([()], lambda x: x.reshape([]))
helper_test_op([(1,)], lambda x: x.reshape([]))
helper_test_op([()], lambda x: x.reshape([1]))
helper_test_op([()], lambda x: x.reshape([1, 1, 1]))
self.helper_test_exception([(3, 4)], lambda x: x.reshape((-1, -1, 2)), lambda x: x.reshape((-1, -1, 2)), expected=RuntimeError)
self.helper_test_exception([(3, 4)], lambda x: x.reshape((-1, -1, -1, 2)), lambda x: x.reshape((-1, -1, -1, 2)), expected=RuntimeError)
helper_test_op([()], lambda x: x.reshape([1,1,1]))
self.helper_test_exception([(3,4)], lambda x: x.reshape((-1,-1,2)), lambda x: x.reshape((-1,-1,2)), expected=RuntimeError)
self.helper_test_exception([(3,4)], lambda x: x.reshape((-1,-1,-1,2)), lambda x: x.reshape((-1,-1,-1,2)), expected=RuntimeError)
with self.assertRaises(ValueError):
x = Tensor.ones((4,3,6,6))
@@ -1139,7 +1139,10 @@ class TestOps(unittest.TestCase):
helper_test_op([(4,3,6,6)], lambda x: x.flip((-1,)))
helper_test_op([()], lambda x: x.flip(()))
helper_test_op([(1,)], lambda x: x.flip(()))
helper_test_op([(4, 3, 6, 6)], lambda x: x.flip(()))
helper_test_op([(4,3,6,6)], lambda x: x.flip(()))
self.helper_test_exception([(3,4)], lambda x: x.flip((0,0)), lambda x: x.flip((0,0)), expected=RuntimeError)
self.helper_test_exception([(3,4)], lambda x: x.flip((1,1)), lambda x: x.flip((1,1)), expected=RuntimeError)
self.helper_test_exception([(3,4)], lambda x: x.flip((1,-1)), lambda x: x.flip((1,-1)), expected=RuntimeError)
def test_squeeze(self):
helper_test_op([(1,3,6,6)], lambda x: x.squeeze(0))
tinygrad/function.py
+1 -1
View File
@@ -205,7 +205,7 @@ class Shrink(Function):
class Flip(Function):
def forward(self, x:LazyBuffer, axis:Tuple[int, ...]) -> LazyBuffer:
self.arg = tuple([-1 if i in set(axis) else 1 for i in range(len(x.shape))])
self.arg = tuple([-1 if i in axis else 1 for i in range(len(x.shape))])
return x.stride(self.arg)
def backward(self, grad_output:LazyBuffer) -> LazyBuffer: return grad_output.stride(self.arg)
tinygrad/tensor.py
+4 -2
View File
@@ -7,7 +7,7 @@ from collections import defaultdict
import numpy as np
from tinygrad.dtype import DType, dtypes, ImageDType, ConstType, least_upper_float, least_upper_dtype, sum_acc_dtype
from tinygrad.helpers import argfix, make_pair, flatten, prod, all_int, round_up, merge_dicts, argsort, getenv, get_shape, fully_flatten
from tinygrad.helpers import argfix, make_pair, flatten, prod, all_int, round_up, merge_dicts, argsort, getenv, get_shape, fully_flatten, dedup
from tinygrad.helpers import IMAGE, DEBUG, WINO, THREEFRY
from tinygrad.lazy import LazyBuffer
from tinygrad.multi import MultiLazyBuffer
@@ -817,7 +817,9 @@ class Tensor:
print(t.flip((0, 1)).numpy())
```
"""
return F.Flip.apply(self, axis=[x if x >= 0 else x+len(self.shape) for x in argfix(axis, *args)])
axis_arg = tuple(self._resolve_dim(x) for x in argfix(axis, *args))
if len(axis_arg) != len(dedup(axis_arg)): raise RuntimeError(f"dim can appear at least once, getting {axis_arg}")
return F.Flip.apply(self, axis=axis_arg)
def shrink(self, arg:Tuple[Optional[Tuple[sint, sint]], ...]) -> Tensor:
"""