import taichi as ti

loss = ti.field(dtype = ti.f32, shape = (), needs_grad = True)
x = ti.field(dtype = ti.f32, shape = (), needs_grad = True)

@ti.kernel
def work():
    loss[None] = x[None] * ti.random()
    # loss[None] = x[None] * 0.5 # this works

x[None] = 10
with ti.Tape(loss):
    work()
print(x.grad[None])