已合并
Skip scatter reduce fallback shim #39744
Skip scatter reduce fallback shim #39744
已合并
fgd_dragon创建于 7月1日
1 个文件变更+3-1
@@ -48,14 +48,16 @@ NPU_UNSUPPORTED_OPS = [
48 "aten._scaled_dot_product_fused_attention_overrideable.default",48 "aten._scaled_dot_product_fused_attention_overrideable.default",
49 "aten._scaled_mm.default",49 "aten._scaled_mm.default",
50 "aten._scaled_grouped_mm.default",50 "aten._scaled_grouped_mm.default",
51+ "aten.scatter_reduce.two",
51 "aten.scatter_reduce.two_out"52 "aten.scatter_reduce.two_out"
52]53]
53 54 
54 55 
56+# Use None as default to avoid FLAKE8 B006 mutable argument warning.
55def _get_backend_index_for_npu(57def _get_backend_index_for_npu(
56 func: NativeFunction,58 func: NativeFunction,
57 backend_indices: dict[DispatchKey, BackendIndex],59 backend_indices: dict[DispatchKey, BackendIndex],
58- structured_func_group_dict: dict[OperatorName, NativeFunctionsGroup] = {},60+ structured_func_group_dict: dict[OperatorName, NativeFunctionsGroup] | None = None,
59) -> BackendIndex | None:61) -> BackendIndex | None:
60 for dk in NPU_DISPATCH_KEYS:62 for dk in NPU_DISPATCH_KEYS:
61 if dk in backend_indices:63 if dk in backend_indices: