-
Notifications
You must be signed in to change notification settings - Fork 3.1k
Expand file tree
/
Copy pathndarray.py
More file actions
312 lines (245 loc) · 6.37 KB
/
Copy pathndarray.py
File metadata and controls
312 lines (245 loc) · 6.37 KB
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
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
243
244
245
246
247
248
249
250
251
252
253
254
255
256
257
258
259
260
261
262
263
264
265
266
267
268
269
270
271
272
273
274
275
276
277
278
279
280
281
282
283
284
285
286
287
288
289
290
291
292
293
294
295
296
297
298
299
300
301
302
303
304
305
306
307
308
309
310
311
312
"""DGL Runtime NDArray API.
dgl.ndarray provides a minimum runtime array structure to be
used with C++ library.
"""
# pylint: disable=invalid-name,unused-import
from __future__ import absolute_import as _abs
import ctypes
import functools
import operator
import numpy as _np
from . import backend as F
from ._ffi.function import _init_api
from ._ffi.ndarray import (
_set_class_ndarray,
context,
DGLContext,
DGLDataType,
empty,
empty_shared_mem,
from_dlpack,
NDArrayBase,
numpyasarray,
)
from ._ffi.object import ObjectBase, register_object
class NDArray(NDArrayBase):
"""Lightweight NDArray class for DGL framework."""
def __len__(self):
return functools.reduce(operator.mul, self.shape, 1)
def shared_memory(self, name):
"""Return a copy of the ndarray in shared memory
Parameters
----------
name : str
The name of the shared memory
Returns
-------
NDArray
"""
return empty_shared_mem(name, True, self.shape, self.dtype).copyfrom(
self
)
def cpu(dev_id=0):
"""Construct a CPU device
Parameters
----------
dev_id : int, optional
The integer device id
Returns
-------
ctx : DGLContext
The created context
"""
return DGLContext(1, dev_id)
def gpu(dev_id=0):
"""Construct a CPU device
Parameters
----------
dev_id : int, optional
The integer device id
Returns
-------
ctx : DGLContext
The created context
"""
return DGLContext(2, dev_id)
def array(arr, ctx=cpu(0)):
"""Create an array from source arr.
Parameters
----------
arr : numpy.ndarray
The array to be copied from
ctx : DGLContext, optional
The device context to create the array
Returns
-------
ret : NDArray
The created array
"""
if not isinstance(arr, (_np.ndarray, NDArray)):
arr = _np.array(arr)
return empty(arr.shape, arr.dtype, ctx).copyfrom(arr)
def zerocopy_from_numpy(np_data):
"""Create an array that shares the given numpy data.
Parameters
----------
np_data : numpy.ndarray
The numpy data
Returns
-------
NDArray
The array
"""
arr, _ = numpyasarray(np_data)
handle = ctypes.pointer(arr)
return NDArray(handle, is_view=True)
def cast_to_signed(arr):
"""Cast this NDArray from unsigned integer to signed one.
uint64 -> int64
uint32 -> int32
Useful for backends with poor signed integer support (e.g., TensorFlow).
Parameters
----------
arr : NDArray
Input array
Returns
-------
NDArray
Cased array
"""
return _CAPI_DGLArrayCastToSigned(arr)
def get_shared_mem_array(name, shape, dtype):
"""Get a tensor from shared memory with specific name
Parameters
----------
name : str
The unique name of the shared memory
shape : tuple of int
The shape of the returned tensor
dtype : F.dtype
The dtype of the returned tensor
Returns
-------
F.tensor
The tensor got from shared memory.
"""
new_arr = empty_shared_mem(
name, False, shape, F.reverse_data_type_dict[dtype]
)
dlpack = new_arr.to_dlpack()
return F.zerocopy_from_dlpack(dlpack)
def create_shared_mem_array(name, shape, dtype):
"""Create a tensor from shared memory with the specific name
Parameters
----------
name : str
The unique name of the shared memory
shape : tuple of int
The shape of the returned tensor
dtype : F.dtype
The dtype of the returned tensor
Returns
-------
F.tensor
The created tensor.
"""
new_arr = empty_shared_mem(
name, True, shape, F.reverse_data_type_dict[dtype]
)
dlpack = new_arr.to_dlpack()
return F.zerocopy_from_dlpack(dlpack)
def exist_shared_mem_array(name):
"""Check the existence of shared-memory array.
Parameters
----------
name : str
The name of the shared-memory array.
Returns
-------
bool
The existence of the array
"""
return _CAPI_DGLExistSharedMemArray(name)
class SparseFormat:
"""Format code"""
ANY = 0
COO = 1
CSR = 2
CSC = 3
FORMAT2STR = {
0: "ANY",
1: "COO",
2: "CSR",
3: "CSC",
}
@register_object("aten.SparseMatrix")
class SparseMatrix(ObjectBase):
"""Sparse matrix object class in C++ backend."""
@property
def format(self):
"""Sparse format enum
Returns
-------
int
"""
return _CAPI_DGLSparseMatrixGetFormat(self)
@property
def num_rows(self):
"""Number of rows.
Returns
-------
int
"""
return _CAPI_DGLSparseMatrixGetNumRows(self)
@property
def num_cols(self):
"""Number of rows.
Returns
-------
int
"""
return _CAPI_DGLSparseMatrixGetNumCols(self)
@property
def indices(self):
"""Index arrays.
Returns
-------
list of ndarrays
"""
ret = [_CAPI_DGLSparseMatrixGetIndices(self, i) for i in range(3)]
return [F.zerocopy_from_dgl_ndarray(arr) for arr in ret]
@property
def flags(self):
"""Flag arrays
Returns
-------
list of boolean
"""
return _CAPI_DGLSparseMatrixGetFlags(self)
def __getstate__(self):
return (
self.format,
self.num_rows,
self.num_cols,
self.indices,
self.flags,
)
def __setstate__(self, state):
fmt, nrows, ncols, indices, flags = state
indices = [F.zerocopy_to_dgl_ndarray(idx) for idx in indices]
self.__init_handle_by_constructor__(
_CAPI_DGLCreateSparseMatrix, fmt, nrows, ncols, indices, flags
)
def __repr__(self):
return 'SparseMatrix(fmt="{}", shape=({},{}))'.format(
SparseFormat.FORMAT2STR[self.format], self.num_rows, self.num_cols
)
_set_class_ndarray(NDArray)
_init_api("dgl.ndarray")
_init_api("dgl.ndarray.uvm", __name__)
# An array representing null (no value) that can be safely converted to
# other backend tensors.
NULL = {
"int64": array(_np.array([], dtype=_np.int64)),
"int32": array(_np.array([], dtype=_np.int32)),
}