Open samuela opened 3 years ago
Thanks! I believe it's because Taichi's Autodiff pass has swallowed this silently... See https://github.com/taichi-dev/taichi/blob/4a568522572fbaaa8ef9c7e594553361e329aaa0/taichi/transforms/auto_diff.cpp#L654
We could check failure here, though I'm not sure if there's any other consideration not to do so.
Adding an even simpler reproducer:
import numpy as np
import taichi as ti
np.random.seed(123)
real = ti.f32
ti.init(ti.cpu, default_fp=real, print_ir=True)
loss = ti.field(dtype=real, shape=(), needs_grad=True)
x = ti.field(dtype=real, shape=(), needs_grad=True)
acc = ti.field(dtype=real, shape=(), needs_grad=False)
@ti.kernel
def forces():
acc[None] = x[None]
@ti.kernel
def vjp():
loss[None] = acc[None]
x[None] = np.random.randn(1).astype(np.float32)[0]
with ti.Tape(loss):
forces()
vjp()
# These should be non-zero!
print(x.grad.to_numpy())
Because acc
's grad is not instantiated, eventually kernel vjp_c7_0_grad_grad
gets optimized to empty...
[I 10/24/20 12:55:09.690] [compile_to_offloads.cpp:operator()@18] [vjp_c7_0_grad_grad] Initial IR:
kernel {
#@tmp0[] = gbl load #@tmp4[]
}
...
[I 10/24/20 12:55:09.691] [compile_to_offloads.cpp:operator()@18] [vjp_c7_0_grad_grad] Simplified IV:
kernel {
}
Describe the bug I've found that it's possible to get Taichi to return incorrect gradients with certain
needs_grad
settings.To Reproduce
Log/Screenshots
even though the correct gradients are non-zero.
Additional comments