How to use the taichi.kernel function in taichi

  @ti.kernel
  def compute_loss():
    total_loss[None] = 0.0
    for i in range(steps):
      total_loss[None].atomic_add(running_loss[i] * 2)
    total_loss[None].atomic_add(additional_loss[None] * 3)

@ti.kernel
def compute_loss(t: ti.i32):
  loss[None] = x[t, head_id][0]

@ti.kernel
def laplace():
  print(5)
  for i, j in x:
    if (i + j) % 3 == 0:
      y[i, j] = 4.0 * x[i, j] - x[i - 1, j] - x[i + 1, j] - x[i, j - 1] - x[i, j + 1]
    else:
      y[i, j] = 0.0

@ti.kernel
def nn1(t: ti.i32):
  for i in range(n_hidden):
    actuation = 0.0
    for j in ti.static(range(n_sin_waves)):
      actuation += weights1[i, j] * ti.sin(spring_omega * t * dt +
                                           2 * math.pi / n_sin_waves * j)
    for j in ti.static(range(n_objects)):
      offset = x[t, j] - center[t]
      # use a smaller weight since there are too many of them
      actuation += weights1[i, j * 4 + n_sin_waves] * offset[0] * 0.05
      actuation += weights1[i, j * 4 + n_sin_waves + 1] * offset[1] * 0.05
      actuation += weights1[i, j * 4 + n_sin_waves + 2] * v[t, i][0] * 0.05
      actuation += weights1[i, j * 4 + n_sin_waves + 3] * v[t, i][1] * 0.05
    actuation += weights1[i, n_objects * 4 + n_sin_waves] * target_v[t][0]
    actuation += weights1[i, n_objects * 4 + n_sin_waves + 1] * target_v[t][1]
    actuation += bias1[i]

@ti.kernel
def apply_grad():
  # gradient descent
  for i in range(n_grid):
    for j in range(n_grid):
      v[0, i, j] -= learning_rate * v.grad[0, i, j]

@ti.kernel
def compute_loss():
  dist = x_avg[None][0]
  loss[None] = -dist

@ti.kernel
def compute_loss(t: ti.i32):
  x01 = x[t, 0] - x[t, 1]
  x02 = x[t, 0] - x[t, 2]
  area = ti.abs(
      0.5 * (x01[0] * x02[1] - x01[1] * x02[0]))  # area from cross product
  target_area = 0.1
  loss[None] = ti.sqr(area - target_area)

@ti.kernel
def advect(field: ti.template(), field_out: ti.template(),
           t_offset: ti.template(), t: ti.i32):
  """Move field smoke according to x and y velocities (vx and vy)
     using an implicit Euler integrator."""
  for y in range(n_grid):
    for x in range(n_grid):
      center_x = y - v[t + t_offset, y, x][0]
      center_y = x - v[t + t_offset, y, x][1]

      # Compute indices of source cell
      left_ix = ti.cast(ti.floor(center_x), ti.i32)
      top_ix = ti.cast(ti.floor(center_y), ti.i32)

      rw = center_x - left_ix  # Relative weight of right-hand cell
      bw = center_y - top_ix  # Relative weight of bottom cell

@ti.kernel
def nn2(t: ti.i32):
  for i in range(n_springs):
    act = 0.0
    for j in ti.static(range(n_hidden)):
      act += weights2[i, j] * hidden[t, j]
    act += bias2[i]
    act = ti.tanh(act)
    actuation[t, i] = act

@ti.kernel
def compute_loss(view_id: ti.i32):
  for i in range(res):
    for j in range(res):
      ti.atomic_add(
          loss,
          ti.sqr(images[view_id, i, j] - target_images[view_id, i, j]) *
          (1.0 / (res * res)))