docarray/docarray/typing/tensor/ndarray.py at main · docarray/docarray · GitHub

History

242 lines (190 loc) · 8.05 KB

Raw

100

101

102

103

104

105

106

107

108

109

110

111

112

113

114

115

116

117

118

119

120

121

122

123

124

125

126

127

128

129

130

131

132

133

134

135

136

137

138

139

140

141

142

143

144

145

146

147

148

149

150

151

152

153

154

155

156

157

158

159

160

161

162

163

164

165

166

167

168

169

170

171

172

173

174

175

176

177

178

179

180

181

182

183

184

185

186

187

188

189

190

191

192

193

194

195

196

197

198

199

200

201

202

203

204

205

206

207

208

209

210

211

212

213

214

215

216

217

218

219

220

221

222

223

224

225

226

227

228

229

230

231

232

233

234

235

236

237

238

239

240

241

242

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()

docarray/docarray/typing/tensor/ndarray.py at main · docarray/docarray · GitHub

Uh oh!

Expand file tree

Breadcrumbs

ndarray.py

Latest commit

History

Breadcrumbs

ndarray.py

File metadata and controls