Stmt ... cannot have operand ... error during AutoDiff

Hi all. I am working on a particle-based fluid simulation. I started with the PBF2D example and made it 3D. Now I want to make it end-to-end differentiable. In my attempts, I came across an error message RuntimeError: [verify.cpp:taichi::lang::IRVerifier::basic_verify@40] stmt 5567 cannot have operand 5558. Does anyone know what this issue is about and how to fix it?

For reference, the offending code is the last line in this function, which is a modified version of the one in PBF2D:

    def update_grid(self, frame: ti.i32):
        for i in range(self.num_particles):
            if self.particle_active[frame,i] == 1:
                cell = self.get_cell(self.positions[frame-1,i] + self.total_pos_delta[i])
                # # ti.Vector doesn't seem to support unpacking yet
                # # but we can directly use int Vectors as indices
                offs = self.grid_num_particles[cell].atomic_add(1)
                self.grid2particles[cell, offs] = i

For context, this kernel loops through every particle in the simulation. For each particle, the kernel finds which cell it belongs to in the background grid. The grid is used for faster neighbor checking.

Here self.num_particles is a constant that is fixed at the start of the simulation. self.particle_active is a boolean tensor keeping track of which particle is active in the simulation. A particle i is active and is included in all calculations for frame k if self.particle_active[k,i]==1. Next, self.get_cell() gets the cell indices for a particle position in the background grid. Finally, the offending last line just adds the particle index i to the array that keeps track of which particle is in which cell. As one last note, the forward simulation works perfectly fine, and the error is thrown when the final gradient is computed with ti.Tape.