faster devectorizer with one line

2026-08-01 20:00:55 +00:00
1 changed files with 1 additions and 2 deletions
@@ -125,8 +125,7 @@ def do_devectorize(b:UOp):
  # broadcasting needs to be already unpacked, Invalid matches any dtype and shape
  if not all(x.shape == b.shape or x.base.is_invalid for x in b.src): return None
  src = []
-  for idx in itertools.product(*[range(x) for x in b.shape]):
-    idx_c = [UOp.const(i) for i in idx]
+  for idx_c in itertools.product(*[[UOp.const(i) for i in range(x)] for x in b.shape]):
    src.append(b.replace(dtype=None, src=tuple(x.base if x.base.is_invalid else x.index(*idx_c) for x in b.src)))
  return UOp.stack(*src).reshape(b.shape) if b.op is not Ops.STORE else UOp.group(*src)