dev(feat): add constraint to treetensor

requirements.txt

@@ -1,4 +1,4 @@
-treevalue>=1.4.1
+treevalue>=1.4.5
 torch>=1.1.0,<=1.12.1
 hbutils>=0.6.13
 numpy

test/common/constraints/test_shape.py

@@ -0,0 +1,245 @@
+import numpy as np
+import pytest
+import torch
+
+from treetensor.common import ShapePrefixConstraint, shape_prefix
+
+
+class NumpyShapePrefixConstraint(ShapePrefixConstraint):
+    __type__ = np.ndarray
+
+
+class TorchShapePrefixConstraint(ShapePrefixConstraint):
+    __type__ = torch.Tensor
+
+
+# noinspection DuplicatedCode
+@pytest.mark.unittest
+class TestCommonConstraintsShape:
+    def test_shape_prefix(self):
+        c1 = shape_prefix(2, 3, 4)
+        assert isinstance(c1, ShapePrefixConstraint)
+        assert c1.prefix == (2, 3, 4)
+        assert repr(c1) == '<ShapePrefixConstraint (2, 3, 4)>'
+
+        assert len(c1) == 3
+        assert c1[0] == 2
+        assert c1[1] == 3
+        assert c1[2] == 4
+        with pytest.raises(IndexError):
+            _ = c1[3]
+        assert c1[-1] == 4
+        assert c1[-2] == 3
+        assert c1[-3] == 2
+        assert c1[1:] == (3, 4)
+
+        c1.validate(np.random.rand(2, 3, 4))
+        c1.validate(np.random.rand(2, 3, 4, 5))
+        with pytest.raises(ValueError):
+            c1.validate(np.random.rand(2, 3))
+        with pytest.raises(ValueError):
+            c1.validate(np.random.rand(2, 3, 3))
+        with pytest.raises(ValueError):
+            c1.validate(np.random.rand(2, 3, 3, 4))
+        with pytest.raises(TypeError):
+            c1.validate([2, 3, 4, 5])
+
+        c1.validate(torch.randn(2, 3, 4))
+        c1.validate(torch.randn(2, 3, 4, 5))
+        with pytest.raises(ValueError):
+            c1.validate(torch.randn(2, 3))
+        with pytest.raises(ValueError):
+            c1.validate(torch.randn(2, 3, 3))
+        with pytest.raises(ValueError):
+            c1.validate(torch.randn(2, 3, 3, 4))
+        with pytest.raises(TypeError):
+            c1.validate([2, 3, 4, 5])
+
+        assert c1 == shape_prefix(2, 3, 4)
+        assert not c1 != shape_prefix(2, 3, 4)
+        assert c1 >= shape_prefix(2, 3, 4)
+        assert c1 <= shape_prefix(2, 3, 4)
+        assert not c1 > shape_prefix(2, 3, 4)
+        assert not c1 < shape_prefix(2, 3, 4)
+
+        assert not c1 == shape_prefix(2, 3)
+        assert c1 != shape_prefix(2, 3)
+        assert c1 >= shape_prefix(2, 3)
+        assert not c1 <= shape_prefix(2, 3)
+        assert c1 > shape_prefix(2, 3)
+        assert not c1 < shape_prefix(2, 3)
+
+        assert not c1 == shape_prefix(2, 3, 4, 5)
+        assert c1 != shape_prefix(2, 3, 4, 5)
+        assert not c1 >= shape_prefix(2, 3, 4, 5)
+        assert c1 <= shape_prefix(2, 3, 4, 5)
+        assert not c1 > shape_prefix(2, 3, 4, 5)
+        assert c1 < shape_prefix(2, 3, 4, 5)
+
+        assert not c1 == shape_prefix(2, 3, 3)
+        assert c1 != shape_prefix(2, 3, 3)
+        assert not c1 >= shape_prefix(2, 3, 3)
+        assert not c1 <= shape_prefix(2, 3, 3)
+        assert not c1 > shape_prefix(2, 3, 3)
+        assert not c1 < shape_prefix(2, 3, 3)
+
+        assert not c1 >= np.ndarray
+        assert not c1 > np.ndarray
+        assert not c1 >= torch.Tensor
+        assert not c1 > torch.Tensor
+
+    def test_shape_prefix_numpy(self):
+        def nsp(*prefix):
+            return shape_prefix(*prefix, type_=NumpyShapePrefixConstraint)
+
+        c1 = nsp(2, 3, 4)
+        assert isinstance(c1, NumpyShapePrefixConstraint)
+        assert c1.prefix == (2, 3, 4)
+        assert repr(c1) == '<NumpyShapePrefixConstraint (2, 3, 4)>'
+
+        c1.validate(np.random.rand(2, 3, 4))
+        c1.validate(np.random.rand(2, 3, 4, 5))
+        with pytest.raises(ValueError):
+            c1.validate(np.random.rand(2, 3))
+        with pytest.raises(ValueError):
+            c1.validate(np.random.rand(2, 3, 3))
+        with pytest.raises(ValueError):
+            c1.validate(np.random.rand(2, 3, 3, 4))
+        with pytest.raises(TypeError):
+            c1.validate([2, 3, 4, 5])
+
+        with pytest.raises(TypeError):
+            c1.validate(torch.randn(2, 3, 4))
+        with pytest.raises(TypeError):
+            c1.validate(torch.randn(2, 3, 4, 5))
+        with pytest.raises(TypeError):
+            c1.validate(torch.randn(2, 3))
+        with pytest.raises(TypeError):
+            c1.validate(torch.randn(2, 3, 3))
+        with pytest.raises(TypeError):
+            c1.validate(torch.randn(2, 3, 3, 4))
+        with pytest.raises(TypeError):
+            c1.validate([2, 3, 4, 5])
+
+        assert c1 == nsp(2, 3, 4)
+        assert not c1 != nsp(2, 3, 4)
+        assert c1 >= nsp(2, 3, 4)
+        assert c1 <= nsp(2, 3, 4)
+        assert not c1 > nsp(2, 3, 4)
+        assert not c1 < nsp(2, 3, 4)
+
+        assert not c1 == nsp(2, 3)
+        assert c1 != nsp(2, 3)
+        assert c1 >= nsp(2, 3)
+        assert not c1 <= nsp(2, 3)
+        assert c1 > nsp(2, 3)
+        assert not c1 < nsp(2, 3)
+
+        assert not c1 == nsp(2, 3, 4, 5)
+        assert c1 != nsp(2, 3, 4, 5)
+        assert not c1 >= nsp(2, 3, 4, 5)
+        assert c1 <= nsp(2, 3, 4, 5)
+        assert not c1 > nsp(2, 3, 4, 5)
+        assert c1 < nsp(2, 3, 4, 5)
+
+        assert not c1 == nsp(2, 3, 3)
+        assert c1 != nsp(2, 3, 3)
+        assert not c1 >= nsp(2, 3, 3)
+        assert not c1 <= nsp(2, 3, 3)
+        assert not c1 > nsp(2, 3, 3)
+        assert not c1 < nsp(2, 3, 3)
+
+        assert not c1 == shape_prefix(2, 3, 4)
+        assert c1 != shape_prefix(2, 3, 4)
+        assert c1 >= shape_prefix(2, 3, 4)
+        assert not c1 <= shape_prefix(2, 3, 4)
+        assert c1 > shape_prefix(2, 3, 4)
+        assert not c1 < shape_prefix(2, 3, 4)
+
+        assert c1 >= np.ndarray
+        assert c1 > np.ndarray
+        assert not c1 >= torch.Tensor
+        assert not c1 > torch.Tensor
+
+    def test_shape_prefix_torch(self):
+        def tsp(*prefix):
+            return shape_prefix(*prefix, type_=TorchShapePrefixConstraint)
+
+        c1 = tsp(2, 3, 4)
+        assert isinstance(c1, TorchShapePrefixConstraint)
+        assert c1.prefix == (2, 3, 4)
+        assert repr(c1) == '<TorchShapePrefixConstraint (2, 3, 4)>'
+
+        with pytest.raises(TypeError):
+            c1.validate(np.random.rand(2, 3, 4))
+        with pytest.raises(TypeError):
+            c1.validate(np.random.rand(2, 3, 4, 5))
+        with pytest.raises(TypeError):
+            c1.validate(np.random.rand(2, 3))
+        with pytest.raises(TypeError):
+            c1.validate(np.random.rand(2, 3, 3))
+        with pytest.raises(TypeError):
+            c1.validate(np.random.rand(2, 3, 3, 4))
+        with pytest.raises(TypeError):
+            c1.validate([2, 3, 4, 5])
+
+        c1.validate(torch.randn(2, 3, 4))
+        c1.validate(torch.randn(2, 3, 4, 5))
+        with pytest.raises(ValueError):
+            c1.validate(torch.randn(2, 3))
+        with pytest.raises(ValueError):
+            c1.validate(torch.randn(2, 3, 3))
+        with pytest.raises(ValueError):
+            c1.validate(torch.randn(2, 3, 3, 4))
+        with pytest.raises(TypeError):
+            c1.validate([2, 3, 4, 5])
+
+        assert c1 == tsp(2, 3, 4)
+        assert not c1 != tsp(2, 3, 4)
+        assert c1 >= tsp(2, 3, 4)
+        assert c1 <= tsp(2, 3, 4)
+        assert not c1 > tsp(2, 3, 4)
+        assert not c1 < tsp(2, 3, 4)
+
+        assert not c1 == tsp(2, 3)
+        assert c1 != tsp(2, 3)
+        assert c1 >= tsp(2, 3)
+        assert not c1 <= tsp(2, 3)
+        assert c1 > tsp(2, 3)
+        assert not c1 < tsp(2, 3)
+
+        assert not c1 == tsp(2, 3, 4, 5)
+        assert c1 != tsp(2, 3, 4, 5)
+        assert not c1 >= tsp(2, 3, 4, 5)
+        assert c1 <= tsp(2, 3, 4, 5)
+        assert not c1 > tsp(2, 3, 4, 5)
+        assert c1 < tsp(2, 3, 4, 5)
+
+        assert not c1 == tsp(2, 3, 3)
+        assert c1 != tsp(2, 3, 3)
+        assert not c1 >= tsp(2, 3, 3)
+        assert not c1 <= tsp(2, 3, 3)
+        assert not c1 > tsp(2, 3, 3)
+        assert not c1 < tsp(2, 3, 3)
+
+        assert not c1 == shape_prefix(2, 3, 4)
+        assert c1 != shape_prefix(2, 3, 4)
+        assert c1 >= shape_prefix(2, 3, 4)
+        assert not c1 <= shape_prefix(2, 3, 4)
+        assert c1 > shape_prefix(2, 3, 4)
+        assert not c1 < shape_prefix(2, 3, 4)
+
+        assert not c1 >= np.ndarray
+        assert not c1 > np.ndarray
+        assert c1 >= torch.Tensor
+        assert c1 > torch.Tensor
+
+    def test_shape_prefix_cross(self):
+        c1 = shape_prefix(2, 3, 4, type_=NumpyShapePrefixConstraint)
+        c2 = shape_prefix(2, 3, 4, type_=TorchShapePrefixConstraint)
+        assert not c1 == c2
+        assert c1 != c2
+        assert not c1 >= c2
+        assert not c1 > c2
+        assert not c1 <= c2
+        assert not c1 < c2

test/torch/constraints/test_shape.py

@@ -0,0 +1,97 @@
+import numpy as np
+import pytest
+import torch
+
+import treetensor.torch as ttorch
+from treetensor.torch import TensorShapePrefixConstraint, shape_prefix
+
+
+# noinspection DuplicatedCode
+@pytest.mark.unittest
+class TestCommonConstraintsShape:
+    def test_shape_prefix(self):
+        c1 = shape_prefix(2, 3, 4)
+        assert isinstance(c1, TensorShapePrefixConstraint)
+        assert c1.prefix == (2, 3, 4)
+        assert repr(c1) == '<TensorShapePrefixConstraint (2, 3, 4)>'
+
+        assert len(c1) == 3
+        assert c1[0] == 2
+        assert c1[1] == 3
+        assert c1[2] == 4
+        with pytest.raises(IndexError):
+            _ = c1[3]
+        assert c1[-1] == 4
+        assert c1[-2] == 3
+        assert c1[-3] == 2
+        assert c1[1:] == (3, 4)
+
+        with pytest.raises(TypeError):
+            c1.validate(np.random.rand(2, 3, 4))
+        with pytest.raises(TypeError):
+            c1.validate(np.random.rand(2, 3, 4, 5))
+        with pytest.raises(TypeError):
+            c1.validate(np.random.rand(2, 3))
+        with pytest.raises(TypeError):
+            c1.validate(np.random.rand(2, 3, 3))
+        with pytest.raises(TypeError):
+            c1.validate(np.random.rand(2, 3, 3, 4))
+        with pytest.raises(TypeError):
+            c1.validate([2, 3, 4, 5])
+
+        c1.validate(torch.randn(2, 3, 4))
+        c1.validate(torch.randn(2, 3, 4, 5))
+        with pytest.raises(ValueError):
+            c1.validate(torch.randn(2, 3))
+        with pytest.raises(ValueError):
+            c1.validate(torch.randn(2, 3, 3))
+        with pytest.raises(ValueError):
+            c1.validate(torch.randn(2, 3, 3, 4))
+        with pytest.raises(TypeError):
+            c1.validate([2, 3, 4, 5])
+
+        assert c1 == shape_prefix(2, 3, 4)
+        assert not c1 != shape_prefix(2, 3, 4)
+        assert c1 >= shape_prefix(2, 3, 4)
+        assert c1 <= shape_prefix(2, 3, 4)
+        assert not c1 > shape_prefix(2, 3, 4)
+        assert not c1 < shape_prefix(2, 3, 4)
+
+        assert not c1 == shape_prefix(2, 3)
+        assert c1 != shape_prefix(2, 3)
+        assert c1 >= shape_prefix(2, 3)
+        assert not c1 <= shape_prefix(2, 3)
+        assert c1 > shape_prefix(2, 3)
+        assert not c1 < shape_prefix(2, 3)
+
+        assert not c1 == shape_prefix(2, 3, 4, 5)
+        assert c1 != shape_prefix(2, 3, 4, 5)
+        assert not c1 >= shape_prefix(2, 3, 4, 5)
+        assert c1 <= shape_prefix(2, 3, 4, 5)
+        assert not c1 > shape_prefix(2, 3, 4, 5)
+        assert c1 < shape_prefix(2, 3, 4, 5)
+
+        assert not c1 == shape_prefix(2, 3, 3)
+        assert c1 != shape_prefix(2, 3, 3)
+        assert not c1 >= shape_prefix(2, 3, 3)
+        assert not c1 <= shape_prefix(2, 3, 3)
+        assert not c1 > shape_prefix(2, 3, 3)
+        assert not c1 < shape_prefix(2, 3, 3)
+
+        assert not c1 >= np.ndarray
+        assert not c1 > np.ndarray
+        assert c1 >= torch.Tensor
+        assert c1 > torch.Tensor
+
+    def test_pshape(self):
+        tt = ttorch.tensor({
+            'a': [[0.8479, 1.0074, 0.2725],
+                  [1.1674, 1.0784, 0.0655]],
+            'b': {'x': [[0.2644, 0.7268, 0.2781, 0.6469],
+                        [2.0015, 0.4448, 0.8814, 1.0063],
+                        [0.1847, 0.5864, 0.4417, 0.2117]]},
+        })
+        assert tt.pshape is None
+
+        tt2 = tt.with_constraints(shape_prefix(2, 3), clear=False)
+        assert tt2.pshape == (2, 3)

treetensor/common/__init__.py

@@ -1,3 +1,4 @@
+from .constraints import *
 from .module import *
 from .object import *
 from .proxy import *

treetensor/common/constraints/__init__.py

@@ -0,0 +1 @@
+from .shape import *