diff --git a/tinygrad/codegen/lowerer.py b/tinygrad/codegen/lowerer.py index d7811f5bcd..663a9b0fa6 100644 --- a/tinygrad/codegen/lowerer.py +++ b/tinygrad/codegen/lowerer.py @@ -15,7 +15,7 @@ def shape_to_idx(s, axis_types, start=0): def get_index(ast:UOp) -> IndexContext: axis_types = ast.arg.axis_types if isinstance(ast.arg, KernelInfo) else () - if len(ast.full_shape) != len(axis_types): + if len(ast.full_shape) != len(axis_types) and ast.st is not None: axis_types = tuple([AxisType.REDUCE if resolve(s != fs) else AxisType.LOOP for s,fs in zip(ast.shape, ast.full_shape)]) return IndexContext(axis_types, [], 0) diff --git a/tinygrad/engine/realize.py b/tinygrad/engine/realize.py index 7d6c2b18df..15e8ddf63c 100644 --- a/tinygrad/engine/realize.py +++ b/tinygrad/engine/realize.py @@ -34,8 +34,9 @@ def get_program(ast:UOp, renderer:Renderer|None=None, opts:list[Opt]|None=None) ast = ast.replace(arg=KernelInfo(opts_to_apply=tuple(opts))) try: uops = full_rewrite(ast, renderer) - except RuntimeError: + except RuntimeError as e: print("***** LINEARIZE FAILURE *****") + print(e) print(f"ast = {ast}") raise assert uops[-1].op is Ops.SINK, "last uop must be sink"