From d57a5557e9c78e41cb2d98d3a1bb511092612353 Mon Sep 17 00:00:00 2001 From: HansBug Date: Fri, 23 Dec 2022 16:24:56 +0800 Subject: [PATCH 1/4] dev(feat): add constraint to treetensor --- requirements.txt | 2 +- treetensor/common/object.py | 4 ++-- treetensor/common/trees.py | 11 +++++------ treetensor/torch/size.py | 4 ++-- treetensor/torch/tensor.py | 4 ++-- 5 files changed, 12 insertions(+), 13 deletions(-) diff --git a/requirements.txt b/requirements.txt index 753e304d6e..aab4873bb3 100644 --- a/requirements.txt +++ b/requirements.txt @@ -1,4 +1,4 @@ -treevalue>=1.4.1 +treevalue>=1.4.2 torch>=1.1.0,<=1.12.1 hbutils>=0.6.13 numpy diff --git a/treetensor/common/object.py b/treetensor/common/object.py index c089cc52d7..572a9bc5b7 100644 --- a/treetensor/common/object.py +++ b/treetensor/common/object.py @@ -20,7 +20,7 @@ class Object(BaseTreeStruct, metaclass=clsmeta(_object, allow_dict=True)): Generic object tree class, used in :py:mod:`treetensor.numpy` and :py:mod:`treetensor.torch`. """ - def __init__(self, data): + def __init__(self, data, constraint=None): """ In :class:`treetensor.common.Object`, object or object tree can be initialized. @@ -37,7 +37,7 @@ def __init__(self, data): └── x --> └── c --> 233 """ - BaseTreeStruct.__init__(self, data) + BaseTreeStruct.__init__(self, data, constraint=constraint) @ireduce(builtins.all, piter=list) @method_treelize() diff --git a/treetensor/common/trees.py b/treetensor/common/trees.py index dc334479f0..09dec29ac0 100644 --- a/treetensor/common/trees.py +++ b/treetensor/common/trees.py @@ -1,3 +1,4 @@ +import inspect from functools import partial from typing import Type @@ -59,15 +60,13 @@ def _mapping_func(_, x): _wrapped_func = func_treelize()(func) class _MetaClass(type): - def __call__(cls, data, *args, **kwargs): - if isinstance(data, TreeStorage): - return type.__call__(cls, data) - elif isinstance(data, cls) and not args and not kwargs: - return data + def __call__(cls, data, *args, constraint=None, **kwargs): + if isinstance(data, TreeStorage) or isinstance(data, cls) and not args and not kwargs: + return type.__call__(cls, data, constraint=constraint) _result = _wrapped_func(data, *args, **kwargs) if isinstance(_result, _TempTreeValue): - return type.__call__(cls, _result) + return type.__call__(cls, _result, constraint=constraint) else: return _result diff --git a/treetensor/torch/size.py b/treetensor/torch/size.py index 5abb0c383e..3c63667e92 100644 --- a/treetensor/torch/size.py +++ b/treetensor/torch/size.py @@ -47,7 +47,7 @@ def _new_func(self, value, *args, **kwargs): # noinspection PyTypeChecker @current_names() class Size(Torch, metaclass=clsmeta(torch.Size, allow_dict=True)): - def __init__(self, data): + def __init__(self, data, constraint=None): """ In :class:`treetensor.torch.Size`, it's similar with the original :class:`torch.Size`. @@ -69,7 +69,7 @@ def __init__(self, data): │ └── x --> torch.Size([3, 4]) └── c --> torch.Size([5]) """ - super(Torch, self).__init__(data) + super(Torch, self).__init__(data, constraint=constraint) @doc_from_base() @ireduce(sum) diff --git a/treetensor/torch/tensor.py b/treetensor/torch/tensor.py index f1545da6b7..9536dc380b 100644 --- a/treetensor/torch/tensor.py +++ b/treetensor/torch/tensor.py @@ -70,7 +70,7 @@ class Tensor(Torch, metaclass=_TensorMeta): )(x) # noinspection PyUnusedLocal - def __init__(self, data, *args, **kwargs): + def __init__(self, data, *args, constraint=None, **kwargs): """ In :class:`treetensor.torch.Tensor`, it's similar but a little bit different with the original :class:`torch.Tensor`. @@ -100,7 +100,7 @@ def __init__(self, data, *args, **kwargs): └── c --> tensor([[ True], [False]]) """ - super(Torch, self).__init__(data) + super(Torch, self).__init__(data, constraint=constraint) @method_treelize(return_type=Object) def __get_attr(self, key): From 174e5ce86a46416fce96f4fa9a2fa31407702725 Mon Sep 17 00:00:00 2001 From: HansBug Date: Mon, 2 Jan 2023 23:11:54 +0800 Subject: [PATCH 2/4] dev(hansbug): add shape prefix constraint class --- test/common/constraints/__init__.py | 0 test/common/constraints/test_shape.py | 234 ++++++++++++++++++++++ treetensor/common/__init__.py | 1 + treetensor/common/constraints/__init__.py | 1 + treetensor/common/constraints/shape.py | 50 +++++ 5 files changed, 286 insertions(+) create mode 100644 test/common/constraints/__init__.py create mode 100644 test/common/constraints/test_shape.py create mode 100644 treetensor/common/constraints/__init__.py create mode 100644 treetensor/common/constraints/shape.py diff --git a/test/common/constraints/__init__.py b/test/common/constraints/__init__.py new file mode 100644 index 0000000000..e69de29bb2 diff --git a/test/common/constraints/test_shape.py b/test/common/constraints/test_shape.py new file mode 100644 index 0000000000..53731e21d8 --- /dev/null +++ b/test/common/constraints/test_shape.py @@ -0,0 +1,234 @@ +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) == '' + + 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) == '' + + 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) == '' + + 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 diff --git a/treetensor/common/__init__.py b/treetensor/common/__init__.py index 04448872d9..6e247154b5 100644 --- a/treetensor/common/__init__.py +++ b/treetensor/common/__init__.py @@ -1,3 +1,4 @@ +from .constraints import * from .module import * from .object import * from .proxy import * diff --git a/treetensor/common/constraints/__init__.py b/treetensor/common/constraints/__init__.py new file mode 100644 index 0000000000..0e346277ec --- /dev/null +++ b/treetensor/common/constraints/__init__.py @@ -0,0 +1 @@ +from .shape import * diff --git a/treetensor/common/constraints/shape.py b/treetensor/common/constraints/shape.py new file mode 100644 index 0000000000..c580a4b157 --- /dev/null +++ b/treetensor/common/constraints/shape.py @@ -0,0 +1,50 @@ +from typing import Type, TypeVar, Optional + +from treevalue.tree import ValueConstraint +from treevalue.tree.tree.constraint import TypeConstraint + +__all__ = [ + 'ShapePrefixConstraint', + 'shape_prefix', +] + + +class ShapePrefixConstraint(ValueConstraint): + __type__: Optional[type] = None + + def __init__(self, *prefix): + ValueConstraint.__init__(self) + self.__prefix = prefix + + @property + def prefix(self): + return self.__prefix + + def _validate_value(self, instance): + if self.__type__ and not isinstance(instance, self.__type__): + raise TypeError(f'Invalid type, {self.__type__.__name__!r} expected but {instance!r} found.') + + if not hasattr(instance, 'shape'): + raise TypeError(f'Shape not found for instance {instance!r}.') + shape = instance.shape + if shape[:len(self.__prefix)] != self.__prefix: + raise ValueError(f'Invalid shape prefix, {self.__prefix!r} expected but {shape!r} found.') + + def _features(self): + return self.__prefix + + def _contains(self, other): + if isinstance(other, ShapePrefixConstraint): + return isinstance(self, type(other)) and self.__prefix[:len(other.__prefix)] == other.__prefix + else: + if self.__type__ and isinstance(other, TypeConstraint): + return issubclass(self.__type__, other.type_) + else: + return False + + +_ShapePrefixType = TypeVar('_ShapePrefixType', bound=ShapePrefixConstraint) + + +def shape_prefix(*args, type_: Type[_ShapePrefixType] = ShapePrefixConstraint) -> _ShapePrefixType: + return type_(*args) From 7d50bf4e47cfe1c313ee5a7c6e6f630e9ba44881 Mon Sep 17 00:00:00 2001 From: HansBug Date: Mon, 2 Jan 2023 23:13:04 +0800 Subject: [PATCH 3/4] dev(hansbug): requires treevalue>=1.4.3 --- requirements.txt | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/requirements.txt b/requirements.txt index aab4873bb3..ea1ff4aee2 100644 --- a/requirements.txt +++ b/requirements.txt @@ -1,4 +1,4 @@ -treevalue>=1.4.2 +treevalue>=1.4.3 torch>=1.1.0,<=1.12.1 hbutils>=0.6.13 numpy From 455fd137bcf0c07e55d40729a1939eeb0d6424ea Mon Sep 17 00:00:00 2001 From: HansBug Date: Mon, 13 Feb 2023 14:02:27 +0800 Subject: [PATCH 4/4] dev(hansbug): add pshape --- requirements.txt | 2 +- test/common/constraints/test_shape.py | 11 +++ test/torch/constraints/__init__.py | 0 test/torch/constraints/test_shape.py | 97 ++++++++++++++++++++++++ treetensor/common/constraints/shape.py | 9 ++- treetensor/torch/__init__.py | 3 + treetensor/torch/constraints/__init__.py | 6 ++ treetensor/torch/constraints/shape.py | 16 ++++ treetensor/torch/tensor.py | 11 +++ 9 files changed, 153 insertions(+), 2 deletions(-) create mode 100644 test/torch/constraints/__init__.py create mode 100644 test/torch/constraints/test_shape.py create mode 100644 treetensor/torch/constraints/__init__.py create mode 100644 treetensor/torch/constraints/shape.py diff --git a/requirements.txt b/requirements.txt index ea1ff4aee2..d3a70cf5fd 100644 --- a/requirements.txt +++ b/requirements.txt @@ -1,4 +1,4 @@ -treevalue>=1.4.3 +treevalue>=1.4.5 torch>=1.1.0,<=1.12.1 hbutils>=0.6.13 numpy diff --git a/test/common/constraints/test_shape.py b/test/common/constraints/test_shape.py index 53731e21d8..202a8b3b16 100644 --- a/test/common/constraints/test_shape.py +++ b/test/common/constraints/test_shape.py @@ -22,6 +22,17 @@ def test_shape_prefix(self): assert c1.prefix == (2, 3, 4) assert repr(c1) == '' + 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): diff --git a/test/torch/constraints/__init__.py b/test/torch/constraints/__init__.py new file mode 100644 index 0000000000..e69de29bb2 diff --git a/test/torch/constraints/test_shape.py b/test/torch/constraints/test_shape.py new file mode 100644 index 0000000000..9906962912 --- /dev/null +++ b/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) == '' + + 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) diff --git a/treetensor/common/constraints/shape.py b/treetensor/common/constraints/shape.py index c580a4b157..a4b62776a6 100644 --- a/treetensor/common/constraints/shape.py +++ b/treetensor/common/constraints/shape.py @@ -1,3 +1,4 @@ +from collections.abc import Sequence from typing import Type, TypeVar, Optional from treevalue.tree import ValueConstraint @@ -9,7 +10,7 @@ ] -class ShapePrefixConstraint(ValueConstraint): +class ShapePrefixConstraint(ValueConstraint, Sequence): __type__: Optional[type] = None def __init__(self, *prefix): @@ -20,6 +21,12 @@ def __init__(self, *prefix): def prefix(self): return self.__prefix + def __getitem__(self, index): + return self.__prefix[index] + + def __len__(self) -> int: + return len(self.__prefix) + def _validate_value(self, instance): if self.__type__ and not isinstance(instance, self.__type__): raise TypeError(f'Invalid type, {self.__type__.__name__!r} expected but {instance!r} found.') diff --git a/treetensor/torch/__init__.py b/treetensor/torch/__init__.py index 65ae740842..af591e4355 100644 --- a/treetensor/torch/__init__.py +++ b/treetensor/torch/__init__.py @@ -5,6 +5,8 @@ import torch +from .constraints import * +from .constraints import __all__ as _constraints_all from .funcs import * from .funcs import __all__ as _funcs_all from .funcs.base import get_func_from_torch @@ -17,6 +19,7 @@ from ..config.meta import __VERSION__ __all__ = [ + *_constraints_all, *_funcs_all, *_size_all, *_tensor_all, diff --git a/treetensor/torch/constraints/__init__.py b/treetensor/torch/constraints/__init__.py new file mode 100644 index 0000000000..93f6eb0e00 --- /dev/null +++ b/treetensor/torch/constraints/__init__.py @@ -0,0 +1,6 @@ +from .shape import * +from .shape import __all__ as _shape_all + +__all__ = [ + *_shape_all +] diff --git a/treetensor/torch/constraints/shape.py b/treetensor/torch/constraints/shape.py new file mode 100644 index 0000000000..ff219d8aae --- /dev/null +++ b/treetensor/torch/constraints/shape.py @@ -0,0 +1,16 @@ +import torch + +from ...common.constraints import ShapePrefixConstraint +from ...common.constraints import shape_prefix as _origin_shape_prefix + +__all__ = [ + 'TensorShapePrefixConstraint', 'shape_prefix', +] + + +class TensorShapePrefixConstraint(ShapePrefixConstraint): + __type__ = torch.Tensor + + +def shape_prefix(*shape): + return _origin_shape_prefix(*shape, type_=TensorShapePrefixConstraint) diff --git a/treetensor/torch/tensor.py b/treetensor/torch/tensor.py index 9536dc380b..10cb81dd86 100644 --- a/treetensor/torch/tensor.py +++ b/treetensor/torch/tensor.py @@ -1,9 +1,12 @@ +from typing import Tuple, Optional + import numpy as np import torch as pytorch from hbutils.reflection import post_process from treevalue import method_treelize, TreeValue, typetrans from .base import Torch, rmreduce, post_reduce, auto_reduce +from .constraints import TensorShapePrefixConstraint from .size import Size from .stream import stream_call from ..common import Object, ireduce, clsmeta, return_self, auto_tree, get_tree_proxy @@ -116,6 +119,14 @@ def _attr_extern(self, name): else: return tree + @property + def pshape(self) -> Optional[Tuple[int, ...]]: + constraint = self.constraint.access_first(TensorShapePrefixConstraint) + if constraint: + return constraint.prefix + else: + return None + @property def torch(self): """