diff --git a/test/test_ops.py b/test/test_ops.py index 7043f57f8c..37298d7db7 100644 --- a/test/test_ops.py +++ b/test/test_ops.py @@ -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)) diff --git a/tinygrad/function.py b/tinygrad/function.py index 36a5c32951..395acc7274 100644 --- a/tinygrad/function.py +++ b/tinygrad/function.py @@ -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) diff --git a/tinygrad/tensor.py b/tinygrad/tensor.py index 66e1f49c7e..d034e93bd9 100644 --- a/tinygrad/tensor.py +++ b/tinygrad/tensor.py @@ -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: """