docarray · samsja · Jan 25, 2023 · Jan 17, 2023 · Jan 17, 2023 · Jan 17, 2023
diff --git a/docarray/typing/tensor/abstract_tensor.py b/docarray/typing/tensor/abstract_tensor.py
@@ -7,6 +7,7 @@
     Dict,
     Generic,
     List,
+    Optional,
     Tuple,
     Type,
     TypeVar,
@@ -63,13 +64,28 @@ def __subclasscheck__(cls, subclass):
     def __instancecheck__(cls, instance):
         is_tensor = isinstance(instance, AbstractTensor)
         if is_tensor:  # custom handling
+            _cls = cast(Type[AbstractTensor], cls)
+            if (
+                _cls.__unparametrizedcls__
+            ):  # This is not None if the tensor is parametrized
+                if (
+                    _cls.get_comp_backend().shape(instance)
+                    != _cls.__docarray_target_shape__
+                ):
+                    return False
+                return any(
+                    issubclass(candidate, _cls.__unparametrizedcls__)
+                    for candidate in type(instance).mro()
+                )
             return any(issubclass(candidate, cls) for candidate in type(instance).mro())
         return super().__instancecheck__(instance)
 
 
 class AbstractTensor(Generic[TTensor, T], AbstractType, ABC):
 
     __parametrized_meta__: type = _ParametrizedMeta
+    __unparametrizedcls__: Optional[type] = None
+    __docarray_target_shape__: Optional[Tuple[int, ...]] = None
     _proto_type_name: str
 
     def _to_node_protobuf(self: T) -> 'NodeProto':
@@ -172,6 +188,7 @@ class _ParametrizedTensor(
             cls,  # type: ignore
             metaclass=cls.__parametrized_meta__,  # type: ignore
         ):
+            __unparametrizedcls__ = cls
             __docarray_target_shape__ = shape
 
             @classmethod

diff --git a/docarray/typing/tensor/ndarray.py b/docarray/typing/tensor/ndarray.py
@@ -121,6 +121,8 @@ def validate(
 
     @classmethod
     def _docarray_from_native(cls: Type[T], value: np.ndarray) -> T:
+        if cls.__unparametrizedcls__:  # This is not None if the tensor is parametrized
+            return value.view(cls.__unparametrizedcls__)
         return value.view(cls)
 
     @classmethod

diff --git a/docarray/typing/tensor/torch_tensor.py b/docarray/typing/tensor/torch_tensor.py
@@ -160,7 +160,10 @@ def _docarray_from_native(cls: Type[T], value: torch.Tensor) -> T:
         :param value: the native torch.Tensor
         :return: a TorchTensor
         """
-        value.__class__ = cls
+        if cls.__unparametrizedcls__:  # This is not None if the tensor is parametrized
+            value.__class__ = cls.__unparametrizedcls__
+        else:
+            value.__class__ = cls
         return cast(T, value)
 
     @classmethod

diff --git a/tests/units/typing/tensor/test_cross_backend.py b/tests/units/typing/tensor/test_cross_backend.py
@@ -0,0 +1,14 @@
+import numpy as np
+from pydantic import parse_obj_as
+
+from docarray.typing import NdArray, TorchTensor
+
+
+def test_coercion_behavior():
+    t_np = parse_obj_as(NdArray[128], np.zeros(128))
+    t_th = parse_obj_as(TorchTensor[128], np.zeros(128))
+
+    assert isinstance(t_np, NdArray[128])
+    assert not isinstance(t_np, TorchTensor[128])
+    assert isinstance(t_th, TorchTensor[128])
+    assert not isinstance(t_th, NdArray[128])
diff --git a/tests/units/typing/tensor/test_np_ops.py b/tests/units/typing/tensor/test_np_ops.py
@@ -0,0 +1,23 @@
+import numpy as np
+
+from docarray import BaseDocument
+from docarray.typing import NdArray
+
+
+def test_tensor_ops():
+    class A(BaseDocument):
+        tensor: NdArray[3, 224, 224]
+
+    class B(BaseDocument):
+        tensor: NdArray[3, 112, 224]
+
+    tensor = A(tensor=np.ones((3, 224, 224))).tensor
+    tensord = A(tensor=np.ones((3, 224, 224))).tensor
+    tensorn = np.zeros((3, 224, 224))
+    tensorhalf = B(tensor=np.ones((3, 112, 224))).tensor
+    tensorfull = np.concatenate([tensorhalf, tensorhalf], axis=1)
+
+    assert type(tensor) == NdArray
+    assert type(tensor + tensord) == NdArray
+    assert type(tensor + tensorn) == NdArray
+    assert type(tensor + tensorfull) == NdArray
diff --git a/tests/units/typing/tensor/test_tensor.py b/tests/units/typing/tensor/test_tensor.py
@@ -138,6 +138,8 @@ def test_parametrized_instance():
     assert isinstance(t, np.ndarray)
 
     assert not isinstance(t, NdArray[256])
+    assert not isinstance(t, NdArray[2, 64])
+    assert not isinstance(t, NdArray[2, 2, 32])
 
 
 def test_parametrized_equality():

diff --git a/tests/units/typing/tensor/test_torch_ops.py b/tests/units/typing/tensor/test_torch_ops.py
@@ -0,0 +1,23 @@
+import torch
+
+from docarray import BaseDocument
+from docarray.typing import TorchTensor
+
+
+def test_tensor_ops():
+    class A(BaseDocument):
+        tensor: TorchTensor[3, 224, 224]
+
+    class B(BaseDocument):
+        tensor: TorchTensor[3, 112, 224]
+
+    tensor = A(tensor=torch.ones(3, 224, 224)).tensor
+    tensord = A(tensor=torch.ones(3, 224, 224)).tensor
+    tensorn = torch.zeros(3, 224, 224)
+    tensorhalf = B(tensor=torch.ones(3, 112, 224)).tensor
+    tensorfull = torch.cat([tensorhalf, tensorhalf], dim=1)
+
+    assert type(tensor) == TorchTensor
+    assert type(tensor + tensord) == TorchTensor
+    assert type(tensor + tensorn) == TorchTensor
+    assert type(tensor + tensorfull) == TorchTensor
diff --git a/tests/units/typing/tensor/test_torch_tensor.py b/tests/units/typing/tensor/test_torch_tensor.py
@@ -90,11 +90,15 @@ def test_parametrized():
 
 @pytest.mark.parametrize('shape', [(3, 224, 224), (224, 224, 3)])
 def test_parameterized_tensor_class_name(shape):
-    tensor = parse_obj_as(TorchTensor[3, 224, 224], torch.zeros(shape))
+    MyTT = TorchTensor[3, 224, 224]
+    tensor = parse_obj_as(MyTT, torch.zeros(shape))
 
-    assert tensor.__class__.__name__ == 'TorchTensor[3, 224, 224]'
-    assert tensor.__class__.__qualname__ == 'TorchTensor[3, 224, 224]'
-    assert f'{tensor[0][0][0]}' == 'TorchTensor[3, 224, 224](0.)'
+    assert MyTT.__name__ == 'TorchTensor[3, 224, 224]'
+    assert MyTT.__qualname__ == 'TorchTensor[3, 224, 224]'
+
+    assert tensor.__class__.__name__ == 'TorchTensor'
+    assert tensor.__class__.__qualname__ == 'TorchTensor'
+    assert f'{tensor[0][0][0]}' == 'TorchTensor(0.)'
 
 
 def test_torch_embedding():
@@ -130,6 +134,8 @@ def test_parametrized_instance():
     assert isinstance(t, torch.Tensor)
 
     assert not isinstance(t, TorchTensor[256])
+    assert not isinstance(t, TorchTensor[2, 128])
+    assert not isinstance(t, TorchTensor[2, 2, 64])
 
 
 def test_parametrized_equality():