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.
Thanks for the report and the repro case!
This is currently still an unsupported feature, but we hope to implement it as some point.
Fixed as of 1.4.0