Skip to content

Commit

Permalink
Cythonize the shape property
Browse files Browse the repository at this point in the history
  • Loading branch information
tqchen committed Dec 20, 2019
1 parent 1fdb94b commit f6fb0ad
Show file tree
Hide file tree
Showing 3 changed files with 11 additions and 4 deletions.
5 changes: 5 additions & 0 deletions python/tvm/_ffi/_ctypes/ndarray.py
Original file line number Diff line number Diff line change
Expand Up @@ -90,6 +90,11 @@ def _copyto(self, target_nd):
check_call(_LIB.TVMArrayCopyFromTo(self.handle, target_nd.handle, None))
return target_nd

@property
def shape(self):
"""Shape of this array"""
return tuple(self.handle.contents.shape[i] for i in range(self.handle.contents.ndim))

def to_dlpack(self):
"""Produce an array from a DLPack Tensor without copying memory
Expand Down
5 changes: 5 additions & 0 deletions python/tvm/_ffi/_cython/ndarray.pxi
Original file line number Diff line number Diff line change
Expand Up @@ -68,6 +68,11 @@ cdef class NDArrayBase:
def __set__(self, value):
self._set_handle(value)

property shape:
"""Shape of this array"""
def __get__(self):
return tuple(self.chandle.shape[i] for i in range(self.chandle.ndim))

def __init__(self, handle, is_view):
self._set_handle(handle)
self.c_is_view = is_view
Expand Down
5 changes: 1 addition & 4 deletions python/tvm/_ffi/ndarray.py
Original file line number Diff line number Diff line change
Expand Up @@ -157,10 +157,6 @@ def from_dlpack(dltensor):

class NDArrayBase(_NDArrayBase):
"""A simple Device/CPU Array object in runtime."""
@property
def shape(self):
"""Shape of this array"""
return tuple(self.handle.contents.shape[i] for i in range(self.handle.contents.ndim))

@property
def dtype(self):
Expand Down Expand Up @@ -240,6 +236,7 @@ def copyfrom(self, source_array):
except:
raise TypeError('array must be an array_like data,' +
'type %s is not supported' % str(type(source_array)))

t = TVMType(self.dtype)
shape, dtype = self.shape, self.dtype
if t.lanes > 1:
Expand Down

0 comments on commit f6fb0ad

Please sign in to comment.