refactor!: pydantic 2 pattern for numpy and pytorch typing

BREAKING CHANGE

Author: samedii

lantern/__init__.py

@@ -1,29 +1,23 @@
-from .numpy_seed import numpy_seed
-from .star import star
-from .functional_base import FunctionalBase
-from .tensor import Tensor
-from .numpy import Numpy
 from .epochs import Epochs
+from .functional_base import FunctionalBase
+from .lambda_module import Lambda
 from .metric import Metric
 from .metric_table import MetricTable
-from .lambda_module import Lambda
 from .module_device import module_device
-from .module_train import module_train, module_eval
+from .module_train import module_eval, module_train
+from .numpy import Numpy
+from .numpy_seed import numpy_seed
+from .progress_bar import ProgressBar
 from .requires_grad import requires_grad, requires_nograd
 from .set_learning_rate import set_learning_rate
 from .set_seeds import set_seeds
+from .star import star
+from .tensor import Tensor
 from .worker_init_fn import worker_init_fn
-from .progress_bar import ProgressBar
 
 try:
     from .early_stopping import EarlyStopping
 except ImportError:
     pass
-from .git_info import git_info
-
-from pkg_resources import get_distribution, DistributionNotFound
 
-try:
-    __version__ = get_distribution("pytorch-lantern").version
-except DistributionNotFound:
-    __version__ = "dev"
+from .git_info import git_info

lantern/functional_base.py

@@ -1,16 +1,12 @@
-from pydantic import BaseModel, Extra
+from pydantic import BaseModel
 
 
 class FunctionalBase(BaseModel):
-    class Config:
-        allow_mutation = False
-        extra = Extra.forbid
-
     def map(self, fn, *args, **kwargs):
         return fn(self, *args, **kwargs)
 
     def replace(self, **kwargs):
-        new_dict = self.dict()
+        new_dict = self.model_dump()
         new_dict.update(**kwargs)
         return type(self)(**new_dict)
 

lantern/numpy.py

@@ -1,24 +1,72 @@
 from __future__ import annotations
 
+from typing import Any, List
+
 import numpy as np
 import torch
+from pydantic import GetCoreSchemaHandler, GetJsonSchemaHandler
+from pydantic.json_schema import JsonSchemaValue
+from pydantic_core import core_schema
+from typing_extensions import Annotated
+
+
+def validate_from_list(values: List) -> np.ndarray:
+    return np.array(values)
+
 
+def validate_from_torch(tensor: torch.Tensor) -> np.ndarray:
+    return tensor.numpy()
+
+
+class Numpy:
+    @classmethod
+    def __get_pydantic_core_schema__(
+        cls,
+        _source_type: Any,
+        _handler: GetCoreSchemaHandler,
+    ) -> core_schema.CoreSchema:
+        from_list_schema = core_schema.chain_schema(
+            [
+                core_schema.list_schema(),
+                core_schema.no_info_plain_validator_function(validate_from_list),
+            ]
+        )
+
+        from_torch_schema = core_schema.chain_schema(
+            [
+                core_schema.is_instance_schema(torch.Tensor),
+                core_schema.no_info_plain_validator_function(validate_from_torch),
+            ]
+        )
+
+        return core_schema.json_or_python_schema(
+            json_schema=from_list_schema,
+            python_schema=core_schema.chain_schema(
+                [
+                    core_schema.union_schema(
+                        [
+                            core_schema.is_instance_schema(np.ndarray),
+                            from_list_schema,
+                            from_torch_schema,
+                        ]
+                    ),
+                    core_schema.no_info_plain_validator_function(cls.validate),
+                ]
+            ),
+            serialization=core_schema.plain_serializer_function_ser_schema(
+                lambda instance: instance.tolist()
+            ),
+        )
 
-class Numpy(np.ndarray):
     @classmethod
-    def __get_validators__(cls):
-        yield cls.validate
+    def __get_pydantic_json_schema__(
+        cls, _core_schema: core_schema.CoreSchema, handler: GetJsonSchemaHandler
+    ) -> JsonSchemaValue:
+        return handler(core_schema.list_schema())
 
     @classmethod
-    def validate(cls, data, config=None, field=None) -> np.ndarray:
-        if isinstance(data, cls):
-            return data.view(np.ndarray)
-        elif isinstance(data, np.ndarray):
-            return data
-        elif isinstance(data, torch.Tensor):
-            return data.numpy()
-        else:
-            return np.array(data)
+    def validate(cls, data, config=None, field=None):
+        return data
 
     @classmethod
     def ndim(cls, ndim) -> Numpy:
@@ -229,9 +277,7 @@ def test_base_model():
     from pytest import raises
 
     class Test(BaseModel):
-        images: Numpy.dims("NCHW")
-
-    Test(images=np.ones((10, 3, 32, 32)))
+        images: Annotated[np.ndarray, Numpy.dims("NCHW")]
 
     with raises(ValueError):
         Test(images=np.ones((10, 3, 32)))
@@ -249,7 +295,7 @@ def test_conversion():
     from pydantic import BaseModel
 
     class Test(BaseModel):
-        numbers: Numpy.dims("N")
+        numbers: Annotated[np.ndarray, Numpy.dims("N")]
 
     Test(numbers=[1.1, 2.1, 3.1])
     Test(numbers=torch.tensor([1.1, 2.1, 3.1]))
@@ -270,15 +316,15 @@ def test_dtype():
     from pytest import raises
 
     class Test(BaseModel):
-        numbers: Numpy.uint8()
+        numbers: Annotated[np.ndarray, Numpy.uint8()]
 
     Test(numbers=[1, 2, 3])
 
     with raises(ValueError):
         Test(numbers=[1.5, 2.2, 3.2])
 
     class TestBool(BaseModel):
-        flags: Numpy.bool()
+        flags: Annotated[np.ndarray, Numpy.bool()]
 
     TestBool(flags=[True, False, True])
 
@@ -291,7 +337,7 @@ def test_from_torch():
     from pydantic import BaseModel
 
     class Test(BaseModel):
-        numbers: Numpy
+        numbers: Annotated[np.ndarray, Numpy]
 
     numbers = torch.tensor([1, 2, 3])
     numpy_numbers = Test(numbers=numbers).numbers
@@ -305,7 +351,7 @@ def test_between():
     from pytest import raises
 
     class Test(BaseModel):
-        numbers: Numpy.between(1, 3.5)
+        numbers: Annotated[np.ndarray, Numpy.between(1, 3.5)]
 
     Test(numbers=[1.5, 2.2, 3.2])
 
@@ -318,7 +364,7 @@ def test_gt():
     from pytest import raises
 
     class Test(BaseModel):
-        numbers: Numpy.gt(1)
+        numbers: Annotated[np.ndarray, Numpy.gt(1)]
 
     Test(numbers=[1.5, 2.2, 3.2])
 

lantern/tensor.py

@@ -1,34 +1,72 @@
 from __future__ import annotations
 
+from typing import Any, List
+
 import numpy as np
 import torch
-import torch._C
+from pydantic import GetCoreSchemaHandler, GetJsonSchemaHandler
+from pydantic.json_schema import JsonSchemaValue
+from pydantic_core import core_schema
+from typing_extensions import Annotated
+
 
+def validate_from_list(values: List) -> torch.Tensor:
+    return torch.tensor(values)
 
-class MetaTensor(torch._C._TensorMeta):
-    def __getitem__(self, validate: Tensor) -> Tensor:
-        return validate
 
+def validate_from_numpy(array: np.ndarray) -> torch.Tensor:
+    return torch.from_numpy(array)
 
-class Tensor(torch.Tensor, metaclass=MetaTensor):
+
+class Tensor:
     @classmethod
-    def Validate(cls) -> Tensor:
-        return cls
+    def __get_pydantic_core_schema__(
+        cls,
+        _source_type: Any,
+        _handler: GetCoreSchemaHandler,
+    ) -> core_schema.CoreSchema:
+        from_list_schema = core_schema.chain_schema(
+            [
+                core_schema.list_schema(),
+                core_schema.no_info_plain_validator_function(validate_from_list),
+            ]
+        )
+
+        from_numpy_schema = core_schema.chain_schema(
+            [
+                core_schema.is_instance_schema(np.ndarray),
+                core_schema.no_info_plain_validator_function(validate_from_numpy),
+            ]
+        )
+
+        return core_schema.json_or_python_schema(
+            json_schema=from_list_schema,
+            python_schema=core_schema.chain_schema(
+                [
+                    core_schema.union_schema(
+                        [
+                            core_schema.is_instance_schema(torch.Tensor),
+                            from_list_schema,
+                            from_numpy_schema,
+                        ]
+                    ),
+                    core_schema.no_info_plain_validator_function(cls.validate),
+                ]
+            ),
+            serialization=core_schema.plain_serializer_function_ser_schema(
+                lambda instance: instance.tolist()
+            ),
+        )
 
     @classmethod
-    def __get_validators__(cls):
-        yield cls.validate
+    def __get_pydantic_json_schema__(
+        cls, _core_schema: core_schema.CoreSchema, handler: GetJsonSchemaHandler
+    ) -> JsonSchemaValue:
+        return handler(core_schema.list_schema())
 
     @classmethod
-    def validate(cls, data, config=None, field=None) -> torch.Tensor:
-        if isinstance(data, cls):
-            return torch.tensor(data)
-        elif isinstance(data, torch.Tensor):
-            return data
-        elif isinstance(data, np.ndarray):
-            return torch.from_numpy(data)
-        else:
-            return torch.as_tensor(data)
+    def validate(cls, data, config=None, field=None):
+        return data
 
     @classmethod
     def ndim(cls, ndim) -> Tensor:
@@ -252,14 +290,11 @@ def bool(cls) -> Tensor:
         return cls.dtype(torch.bool)
 
 
-Validate = Tensor
-
-
 def test_base_model():
     from pydantic import BaseModel
 
     class Test(BaseModel):
-        tensor: Tensor.dims("NCHW").float()
+        tensor: Annotated[torch.Tensor, Tensor.dims("NCHW").float()]
 
     Test(tensor=torch.ones(10, 3, 32, 32))
 
@@ -276,8 +311,8 @@ def test_conversion():
     from pydantic import BaseModel
 
     class Test(BaseModel):
-        numbers: Tensor.dims("N")
-        numbers2: Tensor.dims("N")
+        numbers: Annotated[torch.Tensor, Tensor.dims("N")]
+        numbers2: Annotated[torch.Tensor, Tensor.dims("N")]
 
     Test(
         numbers=[1.1, 2.1, 3.1],
@@ -300,15 +335,15 @@ def test_dtype():
     from pytest import raises
 
     class Test(BaseModel):
-        numbers: Tensor.uint8()
+        numbers: Annotated[torch.Tensor, Tensor.uint8()]
 
     Test(numbers=[1, 2, 3])
 
     with raises(ValueError):
         Test(numbers=[1.5, 2.2, 3.2])
 
     class TestBool(BaseModel):
-        flags: Tensor.bool()
+        flags: Annotated[torch.Tensor, Tensor.bool()]
 
     TestBool(flags=[True, False, True])
 
@@ -320,7 +355,7 @@ def test_device():
     from pydantic import BaseModel
 
     class Test(BaseModel):
-        numbers: Tensor.float().cpu()
+        numbers: Annotated[torch.Tensor, Tensor.float().cpu()]
 
     Test(numbers=[1, 2, 3])
 
@@ -329,7 +364,7 @@ def test_from_numpy():
     from pydantic import BaseModel
 
     class Test(BaseModel):
-        numbers: Tensor
+        numbers: Annotated[torch.Tensor, Tensor]
 
     numbers = np.array([1, 2, 3])
     torch_numbers = Test(numbers=numbers).numbers
@@ -343,7 +378,7 @@ def test_ge():
     from pytest import raises
 
     class Test(BaseModel):
-        numbers: Tensor.ge(0)
+        numbers: Annotated[torch.Tensor, Tensor.ge(0)]
 
     Test(numbers=[1.5, 2.2, 3.2])
 
@@ -356,22 +391,22 @@ def test_ne():
     from pytest import raises
 
     class Test(BaseModel):
-        numbers: Tensor.ne(1)
+        numbers: Annotated[torch.Tensor, Tensor.ne(1)]
 
     Test(numbers=[1.5, 2.2, 3.2])
 
     with raises(ValueError):
         Test(numbers=[1, 2.2, 3.2])
 
 
-def test_alternative_syntax():
+def test_shorthand_syntax():
     from pydantic import BaseModel
     from pytest import raises
 
     class Test(BaseModel):
-        numbers: Tensor[Validate.ne(1)]
+        numbers: Tensor.dims("N").float()
 
-    Test(numbers=[1.5, 2.2, 3.2])
+    Test(numbers=[1.5, 2.2, 3.2]).numbers
 
     with raises(ValueError):
-        Test(numbers=[1, 2.2, 3.2])
+        Test(numbers=[[1, 2.2, 3.2], [1, 2, 3]])