dev(hansbug): try adapt torch2

pyproject.toml

@@ -2,7 +2,8 @@
 requires = [
     "setuptools>=42",
     "wheel",
-    "Cython",
+    "cython>=0.29; platform_system != 'Windows'",
+    "cython>=0.29,<3; platform_system == 'Windows'",
 ]
 
 [tool.cibuildwheel]

requirements.txt

@@ -1,5 +1,5 @@
 enum_tools
-graphviz~=0.17
-dill~=0.3.4
+graphviz>=0.17
+dill>=0.3.4
 click>=7.1.0
 hbutils>=0.9.1

test/tree/integration/test_torch.py

@@ -1,6 +1,7 @@
 from unittest import skipUnless
 
 import pytest
+from hbutils.testing import vpip, OS, vpython
 
 from treevalue import FastTreeValue, register_for_torch
 
@@ -60,3 +61,36 @@ class MyTreeValueX(FastTreeValue):
 
         with pytest.warns(UserWarning):
             register_for_torch(MyTreeValueX)
+
+    @skipUnless(vpip('torch') >= '2.0.0' and OS.linux and vpython < '3.11', 'Torch 2 on linux platform required')
+    def test_torch_compile(self):
+        @torch.compile
+        def foo(x, y, t):
+            z = (x + y * 2000) / (t - 100)
+            return z
+
+        a = torch.randn(3, 4)
+        b = torch.randn(3, 4)
+        c = torch.randn(3, 4)
+        assert torch.isclose(foo(a, b, c), (a + b * 2000) / (c - 100)).all()
+
+    @skipUnless(vpip('torch') >= '2.0.0' and OS.linux and vpython < '3.11', 'Torch 2 on linux platform required')
+    def test_torch_compile_buggy(self):
+        @torch.compile
+        def foox(x, y):
+            z = x + y
+            return z
+
+        x = FastTreeValue({
+            'a': torch.randn(3, 4) + 200,
+            'b': torch.randn(5) - 300,
+        })
+        y = FastTreeValue({
+            'a': torch.rand(4) + 500,
+            'b': torch.randn(4, 5) + 1000,
+        })
+
+        _t_isclose = FastTreeValue.func()(torch.isclose)
+
+        assert _t_isclose(foox(x, y), x + y).all() == \
+               FastTreeValue({'a': torch.tensor(True), 'b': torch.tensor(True)})