Repository: docarray/docarray

Issue:
Bug: Shape info in tensor break forward in some pytorch op
# Context

with `TorchTensor[512]` this break in some case `as_expand`

error 
```
TypeError                                 Traceback (most recent call last)
Cell In[182], line 4
      1 for i, batch in enumerate(loader):
      2     #assert batch.image.tensor.shape == (2, 3, 224, 224)
      3     print(model.encode_vision(batch.image).shape)
----> 4     print(model.encode_text(batch.text))
      6     break

Cell In[176], line 12, in CLIP.encode_text(self, texts)
     11 def encode_text(self, texts: Text) -> TorchTensor:
---> 12     return self.bert(input_ids = texts.tokens.input_ids, attention_mask = texts.tokens.attention_mask)

File ~/.config/JetBrains/PyCharmCE2022.3/scratches/laion/venv/lib/python3.9/site-packages/torch/nn/modules/module.py:1194, in Module._call_impl(self, *input, **kwargs)
   1190 # If we don't have any hooks, we want to skip the rest of the logic in
   1191 # this function, and just call forward.
   1192 if not (self._backward_hooks or self._forward_hooks or self._forward_pre_hooks or _global_backward_hooks
   1193         or _global_forward_hooks or _global_forward_pre_hooks):
-> 1194     return forward_call(*input, **kwargs)
   1195 # Do not call functions when jit is used
   1196 full_backward_hooks, non_full_backward_hooks = [], []

File ~/.config/JetBrains/PyCharmCE2022.3/scratches/laion/venv/lib/python3.9/site-packages/transformers/models/distilbert/modeling_distilbert.py:579, in DistilBertModel.forward(self, input_ids, attention_mask, head_mask, inputs_embeds, output_attentions, output_hidden_states, return_dict)
    577 if inputs_embeds is None:
    578     inputs_embeds = self.embeddings(input_ids)  # (bs, seq_length, dim)
--> 579 return self.transformer(
    580     x=inputs_embeds,
    581     attn_mask=attention_mask,
    582     head_mask=head_mask,
    583     output_attentions=output_attentions,
    584     output_hidden_states=output_hidden_states,
    585     return_dict=return_dict,
    586 )

File ~/.config/JetBrains/PyCharmCE2022.3/scratches/laion/venv/lib/python3.9/site-packages/torch/nn/modules/module.py:1194, in Module._call_impl(self, *input, **kwargs)
   1190 # If we don't have any hooks, we want to skip the rest of the logic in
   1191 # this function, and just call forward.
   1192 if not (self._backward_hooks or self._forward_hooks or self._forward_pre_hooks or _global_backward_hooks
   1193         or _global_forward_hooks or _global_forward_pre_hooks):
-> 1194     return forward_call(*input, **kwargs)
   1195 # Do not call functions when jit is used
   1196 full_backward_hooks, non_full_backward_hooks = [], []

File ~/.config/JetBrains/PyCharmCE2022.3/scratches/laion/venv/lib/python3.9/site-packages/transformers/models/distilbert/modeling_distilbert.py:354, in Transformer.forward(self, x, attn_mask, head_mask, output_attentions, output_hidden_states, return_dict)
    351 if output_hidden_states:
    352     all_hidden_states = all_hidden_states + (hidden_state,)
--> 354 layer_outputs = layer_module(
    355     x=hidden_state, attn_mask=attn_mask, head_mask=head_mask[i], output_attentions=output_attentions
    356 )
    357 hidden_state = layer_outputs[-1]
    359 if output_attentions:

File ~/.config/JetBrains/PyCharmCE2022.3/scratches/laion/venv/lib/python3.9/site-packages/torch/nn/modules/module.py:1194, in Module._call_impl(self, *input, **kwargs)
   1190 # If we don't have any hooks, we want to skip the rest of the logic in
   1191 # this function, and just call forward.
   1192 if not (self._backward_hooks or self._forward_hooks or self._forward_pre_hooks or _global_backward_hooks
   1193         or _global_forward_hooks or _global_forward_pre_hooks):
-> 1194     return forward_call(*input, **kwargs)
   1195 # Do not call functions when jit is used
   1196 full_backward_hooks, non_full_backward_hooks = [], []

File ~/.config/JetBrains/PyCharmCE2022.3/scratches/laion/venv/lib/python3.9/site-packages/transformers/models/distilbert/modeling_distilbert.py:289, in TransformerBlock.forward(self, x, attn_mask, head_mask, output_attentions)
    279 """
    280 Parameters:
    281     x: torch.tensor(bs, seq_length, dim)
   (...)
    286     torch.tensor(bs, seq_length, dim) The output of the transformer block contextualization.
    287 """
    288 # Self-Attention
--> 289 sa_output = self.attention(
    290     query=x,
    291     key=x,
    292     value=x,
    293     mask=attn_mask,
    294     head_mask=head_mask,
    295     output_attentions=output_attentions,
    296 )
    297 if output_attentions:
    298     sa_output, sa_weights = sa_output  # (bs, seq_length, dim), (bs, n_heads, seq_length, seq_length)

File ~/.config/JetBrains/PyCharmCE2022.3/scratches/laion/venv/lib/python3.9/site-packages/torch/nn/modules/module.py:1194, in Module._call_impl(self, *input, **kwargs)
   1190 # If we don't have any hooks, we want to skip the rest of the logic in
   1191 # this function, and just call forward.
   1192 if not (self._backward_hooks or self._forward_hooks or self._forward_pre_hooks or _global_backward_hooks
   1193         or _global_forward_hooks or _global_forward_pre_hooks):
-> 1194     return forward_call(*input, **kwargs)
   1195 # Do not call functions when jit is used
   1196 full_backward_hooks, non_full_backward_hooks = [], []

File ~/.config/JetBrains/PyCharmCE2022.3/scratches/laion/venv/lib/python3.9/site-packages/transformers/models/distilbert/modeling_distilbert.py:215, in MultiHeadSelfAttention.forward(self, query, key, value, mask, head_mask, output_attentions)
    213 q = q / math.sqrt(dim_per_head)  # (bs, n_heads, q_length, dim_per_head)
    214 scores = torch.matmul(q, k.transpose(2, 3))  # (bs, n_heads, q_length, k_length)
--> 215 mask = (mask == 0).view(mask_reshp).expand_as(scores)  # (bs, n_heads, q_length, k_length)
    216 scores = scores.masked_fill(
    217     mask, torch.tensor(torch.finfo(scores.dtype).min)
    218 )  # (bs, n_heads, q_length, k_length)
    220 weights = nn.functional.softmax(scores, dim=-1)  # (bs, n_heads, q_length, k_length)

TypeError: no implementation found for 'torch.Tensor.expand_as' on types that implement __torch_function__: [<class 'docarray.typing.tensor.abstract_tensor.TorchTensor[512]'>, <class 'docarray.typing.tensor.abstract_tensor.TorchTensor[512]'>]
```


Discussion:

Resolving patch:
diff --git a/docarray/typing/tensor/abstract_tensor.py b/docarray/typing/tensor/abstract_tensor.py
--- a/docarray/typing/tensor/abstract_tensor.py
+++ b/docarray/typing/tensor/abstract_tensor.py
@@ -15,9 +15,48 @@
 ShapeT = TypeVar('ShapeT')
 
 
+class _ParametrizedMeta(type):
+    """
+    This metaclass ensures that instance and subclass checks on parametrized Tensors
+    are handled as expected:
+
+    assert issubclass(TorchTensor[128], TorchTensor[128])
+    t = parse_obj_as(TorchTensor[128], torch.zeros(128))
+    assert isinstance(t, TorchTensor[128])
+    etc.
+
+    This special handling is needed because every call to `AbstractTensor.__getitem__`
+    creates a new class on the fly.
+    We want technically distinct but identical classes to be considered equal.
+    """
+
+    def __subclasscheck__(cls, subclass):
+        is_tensor = AbstractTensor in subclass.mro()
+        same_parents = is_tensor and cls.mro()[1:] == subclass.mro()[1:]
+
+        subclass_target_shape = getattr(subclass, '__docarray_target_shape__', False)
+        self_target_shape = getattr(cls, '__docarray_target_shape__', False)
+        same_shape = (
+            same_parents
+            and subclass_target_shape
+            and self_target_shape
+            and subclass_target_shape == self_target_shape
+        )
+
+        if same_shape:
+            return True
+        return super().__subclasscheck__(subclass)
+
+    def __instancecheck__(cls, instance):
+        is_tensor = isinstance(instance, AbstractTensor)
+        if is_tensor:  # custom handling
+            return any(issubclass(candidate, cls) for candidate in type(instance).mro())
+        return super().__instancecheck__(instance)
+
+
 class AbstractTensor(Generic[ShapeT], AbstractType, ABC):
 
-    __parametrized_meta__ = type
+    __parametrized_meta__: type = _ParametrizedMeta
     _PROTO_FIELD_NAME: str
 
     @classmethod
@@ -76,7 +115,7 @@ class _ParametrizedTensor(
             cls,  # type: ignore
             metaclass=cls.__parametrized_meta__,  # type: ignore
         ):
-            _docarray_target_shape = shape
+            __docarray_target_shape__ = shape
 
             @classmethod
             def validate(
@@ -86,7 +125,9 @@ def validate(
                 config: 'BaseConfig',
             ):
                 t = super().validate(value, field, config)
-                return _cls.__docarray_validate_shape__(t, _cls._docarray_target_shape)
+                return _cls.__docarray_validate_shape__(
+                    t, _cls.__docarray_target_shape__
+                )
 
         _ParametrizedTensor.__name__ = f'{cls.__name__}[{shape_str}]'
         _ParametrizedTensor.__qualname__ = f'{cls.__qualname__}[{shape_str}]'
diff --git a/docarray/typing/tensor/ndarray.py b/docarray/typing/tensor/ndarray.py
--- a/docarray/typing/tensor/ndarray.py
+++ b/docarray/typing/tensor/ndarray.py
@@ -23,9 +23,19 @@
     from docarray.computation.numpy_backend import NumpyCompBackend
     from docarray.proto import NdArrayProto, NodeProto
 
+from docarray.base_document.base_node import BaseNode
+
 T = TypeVar('T', bound='NdArray')
 ShapeT = TypeVar('ShapeT')
 
+tensor_base: type = type(BaseNode)
+
+
+# the mypy error suppression below should not be necessary anymore once the following
+# is released in mypy: https://github.com/python/mypy/pull/14135
+class metaNumpy(AbstractTensor.__parametrized_meta__, tensor_base):  # type: ignore
+    pass
+
 
 class NdArray(np.ndarray, AbstractTensor, Generic[ShapeT]):
     """
@@ -71,6 +81,7 @@ class MyDoc(BaseDocument):
     """
 
     _PROTO_FIELD_NAME = 'ndarray'
+    __parametrized_meta__ = metaNumpy
 
     @classmethod
     def __get_validators__(cls):
diff --git a/docarray/typing/tensor/torch_tensor.py b/docarray/typing/tensor/torch_tensor.py
--- a/docarray/typing/tensor/torch_tensor.py
+++ b/docarray/typing/tensor/torch_tensor.py
@@ -19,11 +19,17 @@
 T = TypeVar('T', bound='TorchTensor')
 ShapeT = TypeVar('ShapeT')
 
-torch_base = type(torch.Tensor)  # type: Any
-node_base = type(BaseNode)  # type: Any
+torch_base: type = type(torch.Tensor)
+node_base: type = type(BaseNode)
 
 
-class metaTorchAndNode(torch_base, node_base):
+# the mypy error suppression below should not be necessary anymore once the following
+# is released in mypy: https://github.com/python/mypy/pull/14135
+class metaTorchAndNode(
+    AbstractTensor.__parametrized_meta__,  # type: ignore
+    torch_base,  # type: ignore
+    node_base,  # type: ignore
+):  # type: ignore
     pass
 
 


Test patch:
diff --git a/tests/units/typing/tensor/test_tensor.py b/tests/units/typing/tensor/test_tensor.py
--- a/tests/units/typing/tensor/test_tensor.py
+++ b/tests/units/typing/tensor/test_tensor.py
@@ -90,3 +90,39 @@ def test_np_embedding():
     # illegal shape at class creation time
     with pytest.raises(ValueError):
         parse_obj_as(NdArrayEmbedding[128, 128], np.zeros((128, 128)))
+
+
+def test_parametrized_subclass():
+    c1 = NdArray[128]
+    c2 = NdArray[128]
+    assert issubclass(c1, c2)
+    assert issubclass(c1, NdArray)
+    assert issubclass(c1, np.ndarray)
+
+    assert not issubclass(c1, NdArray[256])
+
+
+def test_parametrized_instance():
+    t = parse_obj_as(NdArray[128], np.zeros(128))
+    assert isinstance(t, NdArray[128])
+    assert isinstance(t, NdArray)
+    assert isinstance(t, np.ndarray)
+
+    assert not isinstance(t, NdArray[256])
+
+
+def test_parametrized_equality():
+    t1 = parse_obj_as(NdArray[128], np.zeros(128))
+    t2 = parse_obj_as(NdArray[128], np.zeros(128))
+    t3 = parse_obj_as(NdArray[256], np.zeros(256))
+    assert (t1 == t2).all()
+    assert not t1 == t3
+
+
+def test_parametrized_operations():
+    t1 = parse_obj_as(NdArray[128], np.zeros(128))
+    t2 = parse_obj_as(NdArray[128], np.zeros(128))
+    t_result = t1 + t2
+    assert isinstance(t_result, np.ndarray)
+    assert isinstance(t_result, NdArray)
+    assert isinstance(t_result, NdArray[128])
diff --git a/tests/units/typing/tensor/test_torch_tensor.py b/tests/units/typing/tensor/test_torch_tensor.py
--- a/tests/units/typing/tensor/test_torch_tensor.py
+++ b/tests/units/typing/tensor/test_torch_tensor.py
@@ -82,3 +82,39 @@ def test_torch_embedding():
     # illegal shape at class creation time
     with pytest.raises(ValueError):
         parse_obj_as(TorchEmbedding[128, 128], torch.zeros(128, 128))
+
+
+def test_parametrized_subclass():
+    c1 = TorchTensor[128]
+    c2 = TorchTensor[128]
+    assert issubclass(c1, c2)
+    assert issubclass(c1, TorchTensor)
+    assert issubclass(c1, torch.Tensor)
+
+    assert not issubclass(c1, TorchTensor[256])
+
+
+def test_parametrized_instance():
+    t = parse_obj_as(TorchTensor[128], torch.zeros(128))
+    assert isinstance(t, TorchTensor[128])
+    assert isinstance(t, TorchTensor)
+    assert isinstance(t, torch.Tensor)
+
+    assert not isinstance(t, TorchTensor[256])
+
+
+def test_parametrized_equality():
+    t1 = parse_obj_as(TorchTensor[128], torch.zeros(128))
+    t2 = parse_obj_as(TorchTensor[128], torch.zeros(128))
+    t3 = parse_obj_as(TorchTensor[256], torch.zeros(256))
+    assert (t1 == t2).all()
+    assert not t1 == t3
+
+
+def test_parametrized_operations():
+    t1 = parse_obj_as(TorchTensor[128], torch.zeros(128))
+    t2 = parse_obj_as(TorchTensor[128], torch.zeros(128))
+    t_result = t1 + t2
+    assert isinstance(t_result, torch.Tensor)
+    assert isinstance(t_result, TorchTensor)
+    assert isinstance(t_result, TorchTensor[128])