lowerer works on rangeifed + print exception (#11925)

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)
 

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"