Cannot use range in a customized function if func_grad is implemented

[Xuan Li]

To reproduce:

import warp as wp

@wp.func
def overload_fn(x: float, y: float):
    for i in range(1):
        pass
    return x * 3.0 + y / 3.0, y**2.5

@wp.func_grad(overload_fn)
def overload_fn_grad(x: float, y: float, adj_ret0: float, adj_ret1: float):
    wp.adjoint[x] += x * adj_ret0 * 42.0 + y * adj_ret1 * 10.0
    wp.adjoint[y] += y * adj_ret1 * 3.0

if __name__ == "__main__":
    wp.init()
    @wp.kernel
    def overload_kernel(x: wp.array(dtype=wp.float32), y: wp.array(dtype=wp.float32)):
        tid = wp.tid()
        overload_fn(x[tid], y[tid])
    
    x = wp.array([1.0, 2.0, 3.0], dtype=wp.float32)
    y = wp.array([4.0, 5.0, 6.0], dtype=wp.float32)
    wp.launch(overload_kernel, inputs=[x, y], dim=x.shape[0])

I got the following error:

  File "xxxxxxx/site-packages/warp/codegen.py", line 1547, in get_unroll_range
    max_unroll = adj.builder.options["max_unroll"]
AttributeError: 'NoneType' object has no attribute 'options'

If I delete overload_fn_grad, then the code works fine.

Manually setting max_unroll to 16 can work around.

[Eric Shi]

Thanks for reporting this!

[Nicolas Capens]

Fixed by 670a448. Thanks Eric!