from typing import TYPE_CHECKING, Any, Generic, List, Tuple, Type, TypeVar, Union, cast
import numpy as np
import orjson
from docarray.base_doc.base_node import BaseNode
from docarray.typing.proto_register import _register_proto
from docarray.typing.tensor.abstract_tensor import AbstractTensor
from docarray.utils._internal.misc import ( # noqa
is_jax_available,
is_tf_available,
is_torch_available,
)
jax_available = is_jax_available()
if jax_available:
import jax.numpy as jnp
from docarray.typing.tensor.jaxarray import JaxArray # noqa: F401
torch_available = is_torch_available()
if torch_available:
import torch
from docarray.typing.tensor.torch_tensor import TorchTensor # noqa: F401
tf_available = is_tf_available()
if tf_available:
import tensorflow as tf # type: ignore
from docarray.typing.tensor.tensorflow_tensor import TensorFlowTensor # noqa: F401
if TYPE_CHECKING:
from docarray.computation.numpy_backend import NumpyCompBackend
from docarray.proto import NdArrayProto
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
@_register_proto(proto_type_name='ndarray')
class NdArray(np.ndarray, AbstractTensor, Generic[ShapeT]):
"""
Subclass of `np.ndarray`, intended for use in a Document.
This enables (de)serialization from/to protobuf and json, data validation,
and coercion from compatible types like `torch.Tensor`.
This type can also be used in a parametrized way, specifying the shape of the array.
---
```python
from docarray import BaseDoc
from docarray.typing import NdArray
import numpy as np
class MyDoc(BaseDoc):
arr: NdArray
image_arr: NdArray[3, 224, 224]
square_crop: NdArray[3, 'x', 'x']
random_image: NdArray[3, ...] # first dimension is fixed, can have arbitrary shape
# create a document with tensors
doc = MyDoc(
arr=np.zeros((128,)),
image_arr=np.zeros((3, 224, 224)),
square_crop=np.zeros((3, 64, 64)),
random_image=np.zeros((3, 128, 256)),
)
assert doc.image_arr.shape == (3, 224, 224)
# automatic shape conversion
doc = MyDoc(
arr=np.zeros((128,)),
image_arr=np.zeros((224, 224, 3)), # will reshape to (3, 224, 224)
square_crop=np.zeros((3, 128, 128)),
random_image=np.zeros((3, 64, 128)),
)
assert doc.image_arr.shape == (3, 224, 224)
# !! The following will raise an error due to shape mismatch !!
from pydantic import ValidationError
try:
doc = MyDoc(
arr=np.zeros((128,)),
image_arr=np.zeros((224, 224)), # this will fail validation
square_crop=np.zeros((3, 128, 64)), # this will also fail validation
random_image=np.zeros((4, 64, 128)), # this will also fail validation
)
except ValidationError as e:
pass
```
---
"""
__parametrized_meta__ = metaNumpy
@classmethod
def _docarray_validate(
cls: Type[T],
value: Union[T, np.ndarray, str, List[Any], Tuple[Any], Any],
) -> T:
if isinstance(value, str):
value = orjson.loads(value)
if isinstance(value, np.ndarray):
return cls._docarray_from_native(value)
elif isinstance(value, NdArray):
return cast(T, value)
elif isinstance(value, AbstractTensor):
return cls._docarray_from_native(value._docarray_to_ndarray())
elif torch_available and isinstance(value, torch.Tensor):
return cls._docarray_from_native(value.detach().cpu().numpy())
elif tf_available and isinstance(value, tf.Tensor):
return cls._docarray_from_native(value.numpy())
elif jax_available and isinstance(value, jnp.ndarray):
return cls._docarray_from_native(value.__array__())
elif isinstance(value, list) or isinstance(value, tuple):
try:
arr_from_list: np.ndarray = np.asarray(value)
return cls._docarray_from_native(arr_from_list)
except Exception:
pass # handled below
try:
arr: np.ndarray = np.ndarray(value)
return cls._docarray_from_native(arr)
except Exception:
pass # handled below
raise ValueError(f'Expected a numpy.ndarray compatible type, got {type(value)}')
@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 cast(T, value.view(cls.__unparametrizedcls__))
return value.view(cls)
def _docarray_to_json_compatible(self) -> np.ndarray:
"""
Convert `NdArray` into a json compatible object
:return: a representation of the tensor compatible with orjson
"""
return self.unwrap()
def unwrap(self) -> np.ndarray:
"""
Return the original ndarray without any memory copy.
The original view rest intact and is still a Document `NdArray`
but the return object is a pure `np.ndarray` but both object share
the same memory layout.
---
```python
from docarray.typing import NdArray
import numpy as np
from pydantic import parse_obj_as
t1 = parse_obj_as(NdArray, np.zeros((3, 224, 224)))
t2 = t1.unwrap()
# here t2 is a pure np.ndarray but t1 is still a Docarray NdArray
# But both share the same underlying memory
```
---
:return: a `numpy.ndarray`
"""
return self.view(np.ndarray)
@classmethod
def from_protobuf(cls: Type[T], pb_msg: 'NdArrayProto') -> 'T':
"""
Read ndarray from a proto msg
:param pb_msg:
:return: a numpy array
"""
source = pb_msg.dense
if source.buffer:
x = np.frombuffer(bytearray(source.buffer), dtype=source.dtype)
return cls._docarray_from_native(x.reshape(source.shape))
elif len(source.shape) > 0:
return cls._docarray_from_native(np.zeros(source.shape))
else:
raise ValueError(f'proto message {pb_msg} cannot be cast to a NdArray')
def to_protobuf(self) -> 'NdArrayProto':
"""
Transform self into a NdArrayProto protobuf message
"""
from docarray.proto import NdArrayProto
nd_proto = NdArrayProto()
nd_proto.dense.buffer = self.tobytes()
nd_proto.dense.ClearField('shape')
nd_proto.dense.shape.extend(list(self.shape))
nd_proto.dense.dtype = self.dtype.str
return nd_proto
@staticmethod
def get_comp_backend() -> 'NumpyCompBackend':
"""Return the computational backend of the tensor"""
from docarray.computation.numpy_backend import NumpyCompBackend
return NumpyCompBackend()
def __class_getitem__(cls, item: Any, *args, **kwargs):
# see here for mypy bug: https://github.com/python/mypy/issues/14123
return AbstractTensor.__class_getitem__.__func__(cls, item) # type: ignore
@classmethod
def _docarray_from_ndarray(cls: Type[T], value: np.ndarray) -> T:
"""Create a `tensor from a numpy array
PS: this function is different from `from_ndarray` because it is private under the docarray namesapce.
This allows us to avoid breaking change if one day we introduce a Tensor backend with a `from_ndarray` method.
"""
return cls._docarray_from_native(value)
def _docarray_to_ndarray(self) -> np.ndarray:
"""Create a `tensor from a numpy array
PS: this function is different from `from_ndarray` because it is private under the docarray namesapce.
This allows us to avoid breaking change if one day we introduce a Tensor backend with a `from_ndarray` method.
"""
return self.unwrap()