From fdf3a572d78e4c8f7921770d4c84ec3363b5da4d Mon Sep 17 00:00:00 2001 From: Kevin Wu Date: Mon, 28 Oct 2024 15:22:32 +0000 Subject: [PATCH] Move masked_fill blocklist to appropriate file --- torch_ttnn/passes/lowering/to_tt_guard.py | 5 +++++ torch_ttnn/passes/lowering/to_tt_guard_autogen.py | 5 ----- 2 files changed, 5 insertions(+), 5 deletions(-) diff --git a/torch_ttnn/passes/lowering/to_tt_guard.py b/torch_ttnn/passes/lowering/to_tt_guard.py index 90d1aa748..73816658b 100644 --- a/torch_ttnn/passes/lowering/to_tt_guard.py +++ b/torch_ttnn/passes/lowering/to_tt_guard.py @@ -67,6 +67,10 @@ ], ] +aten_masked_fill_scalar_blocklist = [ + ["Tensor<[2, 1, 7, 7]> self = ?", "Tensor<[2, 1, 7, 7]> mask = ?", "number value = -3.3895313892515355e+38"], +] + # Need to remove this from the blocklist so that yolos can pass aten_view_default_blocklist.remove(["Tensor<[1, 192, 32, 42]> self = ?", "List[int] size = [1, 192, 1344]"]) @@ -187,6 +191,7 @@ GUARD[torch.ops.aten._to_copy.default] = partial(guard_aten, aten__to_copy_default_blocklist) GUARD[torch.ops.aten.unsqueeze.default] = partial(guard_aten, aten_unsqueeze_default_blocklist) GUARD[torch.ops.aten.squeeze.dim] = partial(guard_aten, aten_squeeze_dim_blocklist) +GUARD[torch.ops.aten.masked_fill.Scalar] = partial(guard_aten, aten_masked_fill_scalar_blocklist) def can_lowering_to_ttnn(node): diff --git a/torch_ttnn/passes/lowering/to_tt_guard_autogen.py b/torch_ttnn/passes/lowering/to_tt_guard_autogen.py index d6163f8cf..dcdb1c3db 100644 --- a/torch_ttnn/passes/lowering/to_tt_guard_autogen.py +++ b/torch_ttnn/passes/lowering/to_tt_guard_autogen.py @@ -617,9 +617,6 @@ ["Tensor<[1, 40, 28, 28]> self = ?", "List[int] size = [1, 40, 28, 28]", "List[int] stride = [31360, 784, 28, 1]"], ] aten_mm_default_blocklist = [["Tensor<[1, 21843]> self = ?", "Tensor<[21843, 768]> mat2 = ?"]] -aten_masked_fill_scalar_blocklist = [ - ["Tensor<[2, 1, 7, 7]> self = ?", "Tensor<[2, 1, 7, 7]> mask = ?", "number value = -3.3895313892515355e+38"], -] def get_inputs(node): @@ -677,7 +674,6 @@ def guard_aten(blocklist, node): torch.ops.aten.native_dropout.default: partial(guard_aten, aten_native_dropout_default_blocklist), torch.ops.aten.new_empty_strided.default: partial(guard_aten, aten_new_empty_strided_default_blocklist), torch.ops.aten.mm.default: partial(guard_aten, aten_mm_default_blocklist), - torch.ops.aten.masked_fill.Scalar: partial(guard_aten, aten_masked_fill_scalar_blocklist), } guard_ops = [ @@ -715,5 +711,4 @@ def guard_aten(blocklist, node): "torch.ops.aten.native_dropout.default", "torch.ops.aten.new_empty_strided.default", "torch.ops.aten.mm.default", - "torch.ops.aten.masked_fill.Scalar", ]