.github/workflows/test.yml

@@ -7,6 +7,7 @@ on:
       - dev/*
       - fix/*
       - test/*
+  workflow_dispatch:
 
 jobs:
   unittest:

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-build.txt

-cython>=0.29
+cython>=0.29; platform_system != 'Windows'
+cython>=0.29,<3; platform_system == 'Windows'
 build>=0.7.0
 auditwheel>=4
\ No newline at end of file

requirements.txt

 enum_tools
-graphviz~=0.17
-dill~=0.3.4
+graphviz>=0.17
+dill>=0.3.4
 click>=7.1.0
 hbutils>=0.9.1
\ No newline at end of file

test/tree/integration/test_torch.py

+from typing import Tuple, Mapping
 from unittest import skipUnless
 
 import pytest
+from hbutils.testing import vpip, OS, vpython
 
 from treevalue import FastTreeValue, register_for_torch
 
@@ -60,3 +62,102 @@ class TestTreeIntegrationTorch:
 
         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)})
+
+    @skipUnless(vpip('torch') >= '2.0.0' and OS.linux and vpython < '3.11', 'Torch 2 on linux platform required')
+    def test_with_module(self):
+        from torch import nn
+
+        class MLP(nn.Module):
+            def __init__(self, in_features: int, out_features: int, layers: Tuple[int, ...] = (1024,)):
+                nn.Module.__init__(self)
+                self.in_features = in_features
+                self.out_features = out_features
+                self.layers = layers
+                ios = [self.in_features, *self.layers, self.out_features]
+                self.mlp = nn.Sequential(
+                    *(
+                        nn.Linear(in_, out_, bias=True)
+                        for in_, out_ in zip(ios[:-1], ios[1:])
+                    )
+                )
+
+            def forward(self, x):
+                return self.mlp(x)
+
+        class MultiHeadMLP(nn.Module):
+            def __init__(self, in_features: int, out_features: Mapping[str, int], layers: Tuple[int, ...] = (1024,)):
+                nn.Module.__init__(self)
+                self.in_features = in_features
+                self.out_features = out_features
+                self.layers = layers
+                _networks = {
+                    o_name: MLP(in_features, o_feat, layers)
+                    for o_name, o_feat in self.out_features.items()
+                }
+                self.mlps = nn.ModuleDict(_networks)
+                self._t_mlps = FastTreeValue(_networks)
+
+            def forward(self, x):
+                return self._t_mlps(x)
+
+        net = MultiHeadMLP(
+            20,
+            {'a': 10, 'b': 20, 'c': 14, 'd': 3},
+        )
+        net = torch.compile(net)
+
+        input1 = torch.randn(3, 20)
+        output1 = net(input1)
+        assert output1.shape == FastTreeValue({
+            'a': torch.Size([3, 10]),
+            'b': torch.Size([3, 20]),
+            'c': torch.Size([3, 14]),
+            'd': torch.Size([3, 3]),
+        })
+
+        input2 = FastTreeValue.func()(torch.randn)(FastTreeValue({
+            'a': (3, 20),
+            'b': (4, 20),
+            'c': (20,),
+            'd': (2, 5, 20),
+        }))
+        output2 = net(input2)
+        assert output2.shape == FastTreeValue({
+            'a': torch.Size([3, 10]),
+            'b': torch.Size([4, 20]),
+            'c': torch.Size([14]),
+            'd': torch.Size([2, 5, 3]),
+        })

treevalue/config/meta.py

@@ -7,7 +7,7 @@ Overview:
 __TITLE__ = "treevalue"
 
 #: Version of this project.
-__VERSION__ = "1.4.11"
+__VERSION__ = "1.4.12"
 
 #: Short description of the project, will be included in ``setup.py``.
 __DESCRIPTION__ = 'A flexible, generalized tree-based data structure.'