NVIDIA / warp

A Python framework for high performance GPU simulation and graphics

Home Page:https://nvidia.github.io/warp/

Geek Repo:Geek Repo

Github PK Tool:Github PK Tool

Autodiff is not supported on element-wise operations

xuan-li opened this issue · comments

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.