[REQ] Gradients for composite type in-place array assignments

[daedalus5]

Description

Add correct gradient propagation for composite type in-place array assignments.

Test

import warp as wp

wp.init()

wp.build.clear_kernel_cache()

@wp.kernel
def vectest(x: wp.array(dtype=wp.vec3), y: wp.array(dtype=float)):
    i = wp.tid()

    x[i].y = y[i]

x = wp.zeros(1, dtype=wp.vec3, requires_grad=True)
y = wp.ones(1, dtype=float, requires_grad=True)

tape = wp.Tape()
with tape:
    wp.launch(kernel=vectest, dim=1, inputs=[x, y])

x.grad = wp.ones_like(x, requires_grad=False)
tape.backward()

# should be [1.0]
print(y.grad)

[JonathanKuelz]

Is there currently a workaround to get these gradients before a new version is going to be released?

[daedalus5]

Yes, you can expand the assignment and write to a new array. So for example, the above test could be re-written:

import warp as wp

wp.init()

wp.build.clear_kernel_cache()

@wp.kernel
def vectest(x: wp.array(dtype=wp.vec3), y: wp.array(dtype=float), z: wp.array(dtype=wp.vec3)):
    i = wp.tid()

    a = x[i]
    z[i] = wp.vec3(a.x, y[i], a.z)

x = wp.zeros(1, dtype=wp.vec3, requires_grad=True)
y = wp.ones(1, dtype=float, requires_grad=True)
z = wp.zeros_like(x)

tape = wp.Tape()
with tape:
    wp.launch(kernel=vectest, dim=1, inputs=[x, y, z])

z.grad = wp.ones_like(x, requires_grad=False)
tape.backward()

# should be [0.0, 1.0, 0.0]
print(z)

# should be [1.0, 0.0, 1.0]
print(x.grad)

# should be [1.0]
print(y.grad)

[JonathanKuelz]

Thank you!

[daedalus5]

@JonathanKuelz This has been added in commit 560ad81

It took a few attempts to get this right. Please let us know if you run into any issues.