SpecDB: Add OutTensor specs for add.Tensor & add.Scalar

Differential Revision: D59402158
pytorch-labs · Jul 5, 2024 · 5fb625d · 5fb625d
1 parent 221e7e7
commit 5fb625d
Showing 1 changed file with 18 additions and 2 deletions.
diff --git a/specdb/db.py b/specdb/db.py
@@ -341,7 +341,14 @@
             ),
         ],
         outspec=[
-            OutArg(ArgType.Tensor),
+            OutArg(
+                ArgType.Tensor,
+                constraints=[
+                    cp.Dtype.Eq(
+                        lambda deps: torch.promote_types(deps[0].dtype, deps[1].dtype)
+                    ),
+                ],
+            )
         ],
     ),
     Spec(
@@ -373,7 +380,16 @@
             ),
         ],
         outspec=[
-            OutArg(ArgType.Tensor),
+            OutArg(
+                ArgType.Tensor,
+                constraints=[
+                    cp.Dtype.Eq(
+                        lambda deps: (
+                            fn.promote_type_with_scalar(deps[0].dtype, deps[1])
+                        )
+                    ),
+                ],
+            )
         ],
     ),
     Spec(