Add float to int cast test.

llvm · Jan 8, 2025 · c6a5ee5 · c6a5ee5
1 parent 1df9e86
commit c6a5ee5
Show file tree

Hide file tree

Showing 2 changed files with 21 additions and 0 deletions.
diff --git a/projects/pt1/e2e_testing/xfail_sets.py b/projects/pt1/e2e_testing/xfail_sets.py
@@ -2749,6 +2749,7 @@
 }
 
 ONNX_XFAIL_SET = {
+    "ToDtypeIntFromFloatModule_basic",
     # This test is expected to time out
     "TimeOutModule_basic",
     # Failure - cast error

diff --git a/projects/pt1/python/torch_mlir_e2e_test/test_suite/type_conversion.py b/projects/pt1/python/torch_mlir_e2e_test/test_suite/type_conversion.py
@@ -274,6 +274,26 @@ def ToDtypeFloatFromIntModule_basic(module, tu: TestUtils):
     module.forward(input)
 
 
+class ToDtypeIntFromFloatModule(torch.nn.Module):
+    def __init__(self):
+        super().__init__()
+
+    @export
+    @annotate_args([None, ([-1, -1], torch.float64, True)])
+    def forward(self, x):
+        return torch.ops.aten.to(
+            x,
+            dtype=torch.int64,
+        )
+
+
+@register_test_case(module_factory=lambda: ToDtypeIntFromFloatModule())
+def ToDtypeIntFromFloatModule_basic(module, tu: TestUtils):
+    input = tu.rand(2, 2, low=-5, high=5)
+    input[1][1] = tu.randint(1, 1) + 0.7
+    module.forward(input)
+
+
 class TypeAsSameModule(torch.nn.Module):
     def __init__(self):
         super().__init__()