Goals:

Have parametrized types for our tensor types, where the tensor shape can be specified directly in the type:

class MyDoc(Document):
    tensor: TorchTensor[3, 224, 224]

doc = MyDoc(tensor=torch.zeros(3, 224, 224))  # works
doc = MyDoc(tensor=torch.zeros(224, 224, 3))  # works by reshaping
doc = MyDoc(tensor=torch.zeros(224))  # fails validation

How it works:

The __class_getitem__() that is reponsible for the [...] syntax is defined in a general manner in AbstractTensor.
This means that concrete tensor types like TorchTensor and NdArray only have to implement __validate_shape__(), the rest works automatically.

Todo:

• TorchTensor
• NdArray
• General solution that applies to all tensor types at once
• Test with (de)serialization
• Test with generic Tensor type -> this does not work, not supported
• Documetn in docstring