NVIDIA / warp

A Python framework for high performance GPU simulation and graphics
https://nvidia.github.io/warp/
Other
1.75k stars 148 forks source link

Autodiff is not supported on element-wise operations #218

Open xuan-li opened 3 weeks ago

xuan-li commented 3 weeks ago

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.

c0d1f1ed commented 2 weeks ago

Thanks for the report and the repro case!

This is currently still an unsupported feature, but we hope to implement it as some point.