Custom Gradient Function Variable Naming

[Clemens Schwarke]

Hi,
When I create a custom gradient function like:

@wp.func_grad(example)
def adj_example(a: float, b: float, adj_ret: float):

	...
	
    if (a < 0.0):
        adj_a = adj_ret

    wp.adjoint[a] += adj_a

the generated code looks like this:

static CUDA_CALLABLE void adj_example(
    float32 var_a,
    float32 var_b,
    float32 & var_adj_ret)
{

	...

    // if (vn < 0.0):
    var_9 = (var_a < var_0);
    if (var_9) {
        // adj_a = adj_ret
        wp::copy(var_10, var_7);  <----------------------------
    }
    // wp.adjoint[a] += adj_a
	var_11 = wp::add(adj_a, adj_a);  <----------------------------
    adj_a = var_11;
}

First, adj_a is substituted with var_10 in the wp::copy call. In the next appearance of adj_a in wp::add, it is not substituted with var_10 anymore. This is very unexpected (at least to me). Am I missing something here?

Thanks in advance!

[Eric Heiden]

Hi Clemens,

This is definitely a bug where the variable names seem to clash with the adj_a variable. We will have this fixed in the next release, but for now you could try a different variable name and clear the cache.

@wp.func_grad(example)
def adj_example(a: float, b: float, adj_ret: float):
    # ...

    temp = 0.0  # use a variable name other than `adj_a` here and initialize with 0
    if a < 0.0:
        temp = adj_ret

    wp.adjoint[a] += temp

[Clemens Schwarke]

Hi Eric,
thanks for your suggestion, it is indeed easy to fix so not an issue for me. Just hard to identify in the first place, so thanks for looking into it!