diff --git a/particle_system.py b/particle_system.py
index 2a060eb..76d9a32 100644
--- a/particle_system.py
+++ b/particle_system.py
@@ -66,12 +66,12 @@ def add_particle(self, p, x, v, density, pressure, material, color):
 
     @ti.kernel
     def add_particles(self, new_particles_num: int,
-                      new_particles_positions: ti.ext_arr(),
-                      new_particles_velocity: ti.ext_arr(),
-                      new_particle_density: ti.ext_arr(),
-                      new_particle_pressure: ti.ext_arr(),
-                      new_particles_material: ti.ext_arr(),
-                      new_particles_color: ti.ext_arr()):
+                      new_particles_positions: ti.types.ndarray(),
+                      new_particles_velocity: ti.types.ndarray(),
+                      new_particle_density: ti.types.ndarray(),
+                      new_particle_pressure: ti.types.ndarray(),
+                      new_particles_material: ti.types.ndarray(),
+                      new_particles_color: ti.types.ndarray()):
         for p in range(self.particle_num[None], self.particle_num[None] + new_particles_num):
             v = ti.Vector.zero(float, self.dim)
             x = ti.Vector.zero(float, self.dim)
@@ -101,7 +101,7 @@ def is_valid_cell(self, cell):
     def allocate_particles_to_grid(self):
         for p in range(self.particle_num[None]):
             cell = self.pos_to_index(self.x[p])
-            offset = self.grid_particles_num[cell].atomic_add(1)
+            offset = ti.atomic_add(self.grid_particles_num[cell], 1)
             self.grid_particles[cell, offset] = p
 
     @ti.kernel
@@ -133,13 +133,13 @@ def initialize_particle_system(self):
         self.search_neighbors()
 
     @ti.kernel
-    def copy_to_numpy_nd(self, np_arr: ti.ext_arr(), src_arr: ti.template()):
+    def copy_to_numpy_nd(self, np_arr: ti.types.ndarray(), src_arr: ti.template()):
         for i in range(self.particle_num[None]):
             for j in ti.static(range(self.dim)):
                 np_arr[i, j] = src_arr[i][j]
 
     @ti.kernel
-    def copy_to_numpy(self, np_arr: ti.ext_arr(), src_arr: ti.template()):
+    def copy_to_numpy(self, np_arr: ti.types.ndarray(), src_arr: ti.template()):
         for i in range(self.particle_num[None]):
             np_arr[i] = src_arr[i]