diff --git a/tinygrad/uop/ops.py b/tinygrad/uop/ops.py index f1b7287335..e6baadba4f 100644 --- a/tinygrad/uop/ops.py +++ b/tinygrad/uop/ops.py @@ -83,7 +83,7 @@ def multirange_str(rngs:Iterable[UOp], color=False, pad=None) -> str: return ret def shape_to_shape_arg(arg:tuple[sint, ...]) -> UOp: - if len(arg) == 0: return UOp(Ops.STACK, dtypes.weakint.vec(0)) + if len(arg) == 0: return UOp(Ops.STACK) elif all_int(arg): return UOp.const(dtypes.weakint.vec(len(arg)), arg) else: return UOp(Ops.STACK, dtypes.weakint.vec(len(arg)), tuple(UOp.const(dtypes.weakint, x) if isinstance(x, int) else x for x in arg)) diff --git a/tinygrad/uop/spec.py b/tinygrad/uop/spec.py index 1cbd396c11..ea73831481 100644 --- a/tinygrad/uop/spec.py +++ b/tinygrad/uop/spec.py @@ -150,8 +150,8 @@ spec_tensor = PatternMatcher([ (UPat({Ops.ADD, Ops.MUL, Ops.CDIV, Ops.FLOORDIV}, dtype=dtypes.weakint), lambda: True), # movement ops - (UPat((Ops.RESHAPE, Ops.EXPAND), src=(UPat(), UPat(dtype=dtypes.weakint))), lambda: True), - (UPat((Ops.PAD, Ops.SHRINK), src=(UPat(), UPat(dtype=dtypes.weakint), UPat(dtype=dtypes.weakint)), name="x"), + (UPat((Ops.RESHAPE, Ops.EXPAND), src=(UPat(), UPat())), lambda: True), + (UPat((Ops.PAD, Ops.SHRINK), src=(UPat(), UPat(), UPat()), name="x"), lambda x: x.src[1].dtype.count == x.src[2].dtype.count), (UPat((Ops.PERMUTE, Ops.FLIP), name="mv", src=(UPat(),)), lambda mv: isinstance(mv.arg, tuple)),