[QUESTION] In the warp.autograd.jacobian function, how to use the integer indices to identify the input_output_mask?

[Taiyuan-Zhang]

Hi,

When using the warp.autograd.jacobian function, I want to use the integer indices to identify the input_output_mask but get an empty jacobian. I also noticed that the jacobian function just uses tape to record the kernel rather than launch the kernel since the outputs are not changed.

The following is an example using names as identifications from the official document:

import warp as wp
import warp.autograd

@wp.kernel
def my_kernel(
    a: wp.array(dtype=float), b: wp.array(dtype=wp.vec3),
    out1: wp.array(dtype=wp.vec2), out2: wp.array(dtype=wp.quat),
):
    tid = wp.tid()
    ai, bi = a[tid], b[tid]
    out1[tid] = wp.vec2(ai * wp.length(bi), -ai * wp.dot(bi, wp.vec3(0.1, 1.0, -0.1)))
    out2[tid] = wp.normalize(wp.quat(ai, bi[0], bi[1], bi[2]))

a = wp.array([2.0, -1.0], dtype=wp.float32, requires_grad=True)
b = wp.array([wp.vec3(3.0, 1.0, 2.0), wp.vec3(-4.0, -1.0, 0.0)], dtype=wp.vec3, requires_grad=True)
out1 = wp.zeros(2, dtype=wp.vec2, requires_grad=True)
out2 = wp.zeros(2, dtype=wp.quat, requires_grad=True)

# compute the Jacobian matrices for all input/output pairs of the kernel using the autodiff engine
jacs = wp.autograd.jacobian(
    my_kernel, dim=len(a), inputs=[a, b], outputs=[out1, out2],
    plot_jacobians=True,
    # select which input/output pairs to compute the Jacobian for
    input_output_mask=[("a", "out1"), ("b", "out2")]
)

But it gives empty jacobians when using integers as identifications:

jacs = wp.autograd.jacobian(
    my_kernel, dim=len(a), inputs=[a, b], outputs=[out1, out2],
    plot_jacobians=True,
    # select which input/output pairs to compute the Jacobian for
    input_output_mask=[(0, 0), (1, 1)]
)

I also take a look at the source code, it seems like there is some missing code in the resovle_arg function?

if outputs is None:
        outputs = []
    if input_output_mask is None:
        input_output_mask = []
    arg_names = [arg.label for arg in kernel.adj.args]

    def resolve_arg(name):
        if isinstance(name, int):
            return name
        return arg_names.index(name)

    input_output_mask = [
        (resolve_arg(input_name), resolve_arg(output_name) - len(inputs))
        for input_name, output_name in input_output_mask
    ]
    input_output_mask = set(input_output_mask)

[eric-heiden]

Hi @Taiyuan-Zhang,

Thank you for bringing this issue to our attention! There is indeed a bug in the resolve_arg mechanism that incorrectly subtracts len(inputs) from the user-provided output indices. We will have a fix for it in the next Warp release. For now, you could try to replace these lines in jacobian and jacobian_fd to work around this issue:

    def resolve_arg(name, offset: int = 0):
        if isinstance(name, int):
            return name
        return arg_names.index(name) + offset

    input_output_mask = [
        (resolve_arg(input_name), resolve_arg(output_name, -len(inputs)))
        for input_name, output_name in input_output_mask
    ]

[shi-eric]

This issue should be fixed with the 1.3.2 release: https://github.com/NVIDIA/warp/releases