diff --git a/test/test_ops.py b/test/test_ops.py index 29882050a8..b3e2b27131 100644 --- a/test/test_ops.py +++ b/test/test_ops.py @@ -118,7 +118,7 @@ class TestOps(unittest.TestCase): with self.assertRaises(ValueError): method((2, -3, 0)) def test_negative_dims_full(self): - with self.assertRaises(ValueError): Tensor.full(-3, 2) + with self.assertRaises(ValueError): Tensor.full((-3,), 2) with self.assertRaises(ValueError): Tensor.full((2, -3), 4) with self.assertRaises(ValueError): Tensor.full((2, -3, 0), 4) diff --git a/tinygrad/ops.py b/tinygrad/ops.py index 38ede668d5..e917925b07 100644 --- a/tinygrad/ops.py +++ b/tinygrad/ops.py @@ -40,7 +40,7 @@ T = TypeVar("T") class MathTrait: # required to implement def alu(self:T, arg:Union[UnaryOps, BinaryOps, TernaryOps], *src) -> T: raise NotImplementedError - def const_like(self, b:ConstType|Variable): raise NotImplementedError + def const_like(self, b:ConstType|Variable|Tuple[ConstType]): raise NotImplementedError # great functions you get! def ufix(self, x): return self.const_like(x) if not isinstance(x, MathTrait) else x @@ -384,7 +384,7 @@ class UOp(MathTrait): return ret.arg def sink(self, *srcs): return UOp(UOps.SINK, dtypes.void, (self,)+srcs) def swizzle(self, st:ShapeTracker): return UOp(UOps.SWIZZLE, self.dtype, (self,), st) - def const_like(self, b:ConstType|Variable): return type(self).const(self.dtype, b) + def const_like(self, b:ConstType|Variable|Tuple[ConstType]): return type(self).const(self.dtype, b) def cast(self, dtype:DType): return type(self)(UOps.CAST, dtype, (self,)) def bitcast(self, dtype:DType): return type(self)(UOps.BITCAST, dtype, (self,)) def gep(self, i:Union[Tuple[int, ...], int]): @@ -671,7 +671,7 @@ class UPat(MathTrait): @classmethod def store(cls, *src:UPat): return cls(UOps.STORE, dtypes.void, src) - def const_like(self, b:ConstType|Variable): return type(self).const(self.dtype, b) + def const_like(self, b:ConstType|Variable|Tuple[ConstType]): return type(self).const(self.dtype, b) def alu(self, arg, *src:UPat): asrc = (self,)+src return type(self)(UOps.ALU, None if arg in {BinaryOps.CMPLT, BinaryOps.CMPNE} else asrc[-1].dtype, diff --git a/tinygrad/tensor.py b/tinygrad/tensor.py index 6797b45e53..7d2954dafe 100644 --- a/tinygrad/tensor.py +++ b/tinygrad/tensor.py @@ -1778,7 +1778,7 @@ class Tensor: out_ellipse + ''.join(sorted(c for c in inputs_str if inputs_str.count(c) == 1 and c.isalpha() and c not in out_ellipse))) return formula.split("->") if "->" in formula else (formula, ''.join(c for c in sorted(formula) if formula.count(c) == 1 and c.isalpha())) - xs:Tuple[Tensor] = argfix(*raw_xs) + xs:Tuple[Tensor, ...] = argfix(*raw_xs) inputs_str, output = parse_formula(formula.replace(" ", ""), *xs) inputs = inputs_str.split(",") assert len(xs) == len(inputs), f"number of inputs doesn't match number of operands in formula, expected {len(inputs)}, got {len(xs)}" @@ -1826,7 +1826,7 @@ class Tensor: xup = xup.shrink(tuple(noop_ + flatten(((0,o), (0,k)) for o,k in zip(o_, k_)))) return xup.permute(*range(len(noop_)), *[len(noop_)+i*2 for i in range(len(i_))], *[len(noop_)+i*2+1 for i in range(len(i_))]) - def _padding2d(self, padding:Union[int, Tuple[int, ...]], dims:int) -> Sequence[int]: + def _padding2d(self, padding:Union[int, Sequence[int]], dims:int) -> Sequence[int]: return [padding]*2*dims if isinstance(padding, int) else (padding if len(padding) == 2*dims else [p for p in padding for _ in range(2)][::-1]) # NOTE: these work for more than 2D