add tribal knowledge decomps for newer torch versions

Once we migrate default torch version these will be hit by existing tests.
nod-ai · Feb 6, 2024 · 42821b3 · 42821b3
1 parent 66f79ab
commit 42821b3
Showing 1 changed file with 5 additions and 0 deletions.
diff --git a/core/shark_turbine/dynamo/passes.py b/core/shark_turbine/dynamo/passes.py
@@ -52,6 +52,11 @@
     torch.ops.aten._scaled_dot_product_flash_attention.default,
 ]
 
+# These decompositions either didnt exist or weren't required for 2.1.0
+if torch.__version__ > "2.1.0":
+    DEFAULT_DECOMPOSITIONS.append(torch.ops.aten._scaled_dot_product_flash_attention_for_cpu)
+    DEFAULT_DECOMPOSITIONS.append(torch.ops.aten.unbind_int)
+
 
 def apply_decompositions(
     gm: torch.fx.GraphModule,