diff --git a/tinygrad/schedule/multi.py b/tinygrad/schedule/multi.py index a1cfe18c1a..b6623e0f14 100644 --- a/tinygrad/schedule/multi.py +++ b/tinygrad/schedule/multi.py @@ -241,4 +241,4 @@ multi_pm = PatternMatcher([ ])+replace_allreduce @track_rewrites() -def get_multi_map(big_sink:UOp) -> dict[UOp, UOp]: return graph_rewrite_map(big_sink, multi_pm) +def get_multi_map(big_sink:UOp) -> dict[UOp, UOp]: return graph_rewrite_map(big_sink, multi_pm, name="multi_pm")