[FxImporter] Fix primitive type in return (llvm#3379)

vinayakdsci · May 23, 2024 · d924d00 · d924d00
1 parent 2e194e1
commit d924d00
Show file tree

Hide file tree

Showing 2 changed files with 7 additions and 1 deletion.
diff --git a/projects/pt1/e2e_testing/xfail_sets.py b/projects/pt1/e2e_testing/xfail_sets.py
@@ -450,7 +450,6 @@
     "TensorToBool_basic",
     "TensorToFloatZeroRank_basic",
     "TensorToFloat_basic",
-    "TestMultipleTensorAndPrimitiveTypesReturn_basic",
     "ThresholdBackward2dMixedModule_basic",
     "TorchPrimLoopForLikeModule_basic",
     "TorchPrimLoopWhileLikeModule_basic",

diff --git a/python/torch_mlir/extras/fx_importer.py b/python/torch_mlir/extras/fx_importer.py
@@ -848,6 +848,13 @@ def _graph_to_function_meta(self, g: Graph) -> Tuple[FunctionType, Location]:
                         result_types.append(
                             self._cc.tensor_to_vtensor_type(result_node)
                         )
+                    elif type(result_node) in SCALAR_TYPE_TO_TORCH_MLIR_TYPE:
+                        result_types.append(
+                            IrType.parse(
+                                SCALAR_TYPE_TO_TORCH_MLIR_TYPE[type(result_node)],
+                                self._c,
+                            )
+                        )
                     else:
                         result_types.append(self._cc.node_val_to_type(result_node))
         return (