[Lang] Add @python_scope decorator for selected MatrixNdarray/VectorN…

…darray methods
taichi-dev · Aug 23, 2022 · 2050072 · 2050072
1 parent 1b4006f
commit 2050072
Show file tree

Hide file tree

Showing 2 changed files with 12 additions and 0 deletions.
diff --git a/python/taichi/lang/_ndarray.py b/python/taichi/lang/_ndarray.py
@@ -75,6 +75,7 @@ def fill(self, val):
         else:
             self._fill_by_kernel(val)
 
+    @python_scope
     def _ndarray_to_numpy(self):
         """Converts ndarray to a numpy array.
 
@@ -88,6 +89,7 @@ def _ndarray_to_numpy(self):
         impl.get_runtime().sync()
         return arr
 
+    @python_scope
     def _ndarray_matrix_to_numpy(self, layout, as_vector):
         """Converts matrix ndarray to a numpy array.
 
@@ -103,6 +105,7 @@ def _ndarray_matrix_to_numpy(self, layout, as_vector):
         impl.get_runtime().sync()
         return arr
 
+    @python_scope
     def _ndarray_from_numpy(self, arr):
         """Loads all values from a numpy array.
 
@@ -122,6 +125,7 @@ def _ndarray_from_numpy(self, arr):
         ext_arr_to_ndarray(arr, self)
         impl.get_runtime().sync()
 
+    @python_scope
     def _ndarray_matrix_from_numpy(self, arr, layout, as_vector):
         """Loads all values from a numpy array.
 
@@ -192,6 +196,7 @@ def _fill_by_kernel(self, val):
         """
         raise NotImplementedError()
 
+    @python_scope
     def _pad_key(self, key):
         if key is None:
             key = ()
@@ -200,6 +205,7 @@ def _pad_key(self, key):
         assert len(key) == len(self.arr.total_shape())
         return key
 
+    @python_scope
     def _initialize_host_accessor(self):
         if self.host_accessor:
             return

diff --git a/python/taichi/lang/matrix.py b/python/taichi/lang/matrix.py
@@ -1957,17 +1957,20 @@ def from_numpy(self, arr):
         """
         self._ndarray_matrix_from_numpy(arr, self.layout, as_vector=0)
 
+    @python_scope
     def __deepcopy__(self, memo=None):
         ret_arr = MatrixNdarray(self.n, self.m, self.dtype, self.shape,
                                 self.layout)
         ret_arr.copy_from(self)
         return ret_arr
 
+    @python_scope
     def _fill_by_kernel(self, val):
         from taichi._kernels import \
             fill_ndarray_matrix  # pylint: disable=C0415
         fill_ndarray_matrix(self, val)
 
+    @python_scope
     def __repr__(self):
         return f'<{self.n}x{self.m} {self.layout} ti.Matrix.ndarray>'
 
@@ -2053,16 +2056,19 @@ def from_numpy(self, arr):
         """
         self._ndarray_matrix_from_numpy(arr, self.layout, as_vector=1)
 
+    @python_scope
     def __deepcopy__(self, memo=None):
         ret_arr = VectorNdarray(self.n, self.dtype, self.shape, self.layout)
         ret_arr.copy_from(self)
         return ret_arr
 
+    @python_scope
     def _fill_by_kernel(self, val):
         from taichi._kernels import \
             fill_ndarray_matrix  # pylint: disable=C0415
         fill_ndarray_matrix(self, val)
 
+    @python_scope
     def __repr__(self):
         return f'<{self.n} {self.layout} ti.Vector.ndarray>'