[QUESTION] How to calculate the gradients of different parameters from the same kernel #407
Description
I am trying to calculate the gradients by using tape = wp.Tape() and tape.backward. Everything works just right when I only call the tape.backward once. However, when I do several calculations in a single kernel and I try to calculate the gradients of several different parameters, the does NOT work as usual.
Here is an example.
import warp as wp
import numpy as np
@wp.kernel
def double(x: wp.array(dtype=float), y: wp.array(dtype=float), z: wp.array(dtype=float), s: wp.array(dtype=float)):
tid = wp.tid()
y[tid] = x[tid]
z[tid] = y[tid]
s[tid] = z[tid]
x = wp.array(np.arange(3), dtype=float, requires_grad=True)
y = wp.zeros_like(x)
z = wp.zeros_like(x)
s = wp.zeros_like(x)
tape = wp.Tape()
with tape:
wp.launch(double, dim=3, inputs=[x, y, z, s])
tape.backward(grads={y: wp.ones_like(x)})
print(x.grad)
tape.zero()
tape.backward(grads={z: wp.ones_like(x)})
print(x.grad)
tape.zero()
tape.backward(grads={s: wp.ones_like(x)})
print(x.grad)
After running this code, it prints
[1. 1. 1.]
[2. 2. 2.]
[4. 4. 4.].
But the gradients of y, z and s are supposed to be [1. 1. 1.]. So my main question is how to calculate the gradients of different parameters from the same kernel?