warp
warp copied to clipboard
Autodiff is not supported on element-wise operations
To reproduce:
import warp as wp
wp.init()
@wp.kernel
def test_grad(a: wp.array(dtype=wp.vec3), b: wp.array(dtype=wp.vec3)):
tid = wp.tid()
ai = a[tid]
# not working
tmp = wp.vec3(0., 0., 0.)
tmp[0] += ai[0]
tmp[1] += ai[1]
tmp[2] += ai[2]
# working
# tmp += ai
b[tid] = tmp
a = wp.ones(3, dtype=wp.vec3, requires_grad=True)
b = wp.ones(3, dtype=wp.vec3, requires_grad=True)
b.grad.fill_(1.)
wp.launch(test_grad, a.shape[0], inputs=[a, b], adjoint=True, adj_inputs=[None, None])
print(a.grad.numpy())
The result is:
[[0. 0. 0.]
[0. 0. 0.]
[0. 0. 0.]]
I am not sure whether this is a bug or an unsupported feature.