Fix nparray >= smth

lezcano · lezcano · commit 0411324bbd42 · 2023-08-08T20:27:00.000Z
diff --git a/test/dynamo/test_misc.py b/test/dynamo/test_misc.py
@@ -1369,6 +1369,17 @@ def fn(y):
         self.assertEqual(r, x * 3)
         del x
 
+    def test_numpy_gt(self):
+        x = np.arange(10)
+
+        @torch.compile
+        def fn(y):
+            return y >= 3
+
+        r = fn(x)
+        self.assertEqual(type(r), np.ndarray)
+        self.assertEqual(r, x >= 3)
+
     def test_graph_break_correctly_when_passing_numpy_ndarray_to_torch_function(self):
         # from transformers/models/big_bird/modeling_big_bird.py
         def fn(x: int, y: torch.Tensor):
diff --git a/torch/_dynamo/variables/builtin.py b/torch/_dynamo/variables/builtin.py
@@ -1339,11 +1339,12 @@ def _unimplemented():
             return ConstantVariable(op(left._underlying_items, right._underlying_items))
 
         if isinstance(left, TensorVariable):
-            from .builder import wrap_fx_proxy
+            from .builder import wrap_fx_proxy_cls
 
             if op not in supported_tensor_comparison_ops.values():
                 _unimplemented()
-            return wrap_fx_proxy(
+            return wrap_fx_proxy_cls(
+                type(left),  # handle Ndarrays and Tensors
                 tx,
                 op(left.as_proxy(), right.as_proxy()),
             )
