From 6f4cfd98faac14a196ddd2654e9835f4ba785f3e Mon Sep 17 00:00:00 2001 From: Yidi Wu Date: Thu, 1 May 2025 16:41:07 -0700 Subject: [PATCH] Avoid directly calling the map operator (#10638) Summary: Calling the higher order operator directly is not recommended since it skips the safety check. Reviewed By: angelayi Differential Revision: D74036464 --- exir/tests/control_flow_models.py | 5 +++-- 1 file changed, 3 insertions(+), 2 deletions(-) diff --git a/exir/tests/control_flow_models.py b/exir/tests/control_flow_models.py index 0815f8fbc5b..5aab85cc45a 100644 --- a/exir/tests/control_flow_models.py +++ b/exir/tests/control_flow_models.py @@ -5,6 +5,7 @@ # LICENSE file in the root directory of this source tree. import torch +from torch._higher_order_ops.map import map as torch_map from torch.nn import Module # @manual @@ -87,7 +88,7 @@ def forward(self, xs, y): def f(x, y): return x + y - return torch.ops.higher_order.map(f, xs, y) + xs + return torch_map(f, xs, y) + xs def get_random_inputs(self): return torch.rand(2, 4), torch.rand(4) @@ -101,7 +102,7 @@ def forward(self, xs, y): def f(x, y): return x + y - return torch.ops.higher_order.map(f, xs, y) + xs + return torch_map(f, xs, y) + xs def get_upper_bound_inputs(self): return torch.rand(4, 4), torch.rand(4)