hotfix: add ctx to VIZ rewrite (#8667)

This commit is contained in:
qazal
2025-01-18 07:58:47 -05:00
committed by GitHub
parent 0ef85b52e6
commit bd0fb14d70

View File

@@ -522,7 +522,7 @@ remove_movement_ops = PatternMatcher([
def create_schedule_with_vars(outs:list[UOp], skip_check:bool=not __debug__) -> tuple[list[ScheduleItem], dict[Variable, int], dict[UOp, UOp]]:
big_sink = UOp.sink(*outs)
# if using VIZ, do a graph rewrite to vizualize the Tensor graph
if getenv("VIZ"): graph_rewrite(big_sink, remove_movement_ops+ops_folding)
if getenv("VIZ"): graph_rewrite(big_sink, remove_movement_ops+ops_folding, ScheduleContext())
if not skip_check: type_verify(list(big_sink.toposort), tensor_uop_spec)
# to_uop is removing (many) of the movement ops
sink = add_buffers(big_sink, ctx:=ScheduleContext(), cache={})