AttributeError: 'FunctionMetadata' object has no attribute 'get' in wp.autograd.jacobian_plot when tracing a Python function

[mathisg-ui]

Bug Description

Bug Description

When using wp.autograd.jacobian with a Python function pipeline (instead of a single @wp.kernel) and setting plot_jacobians=True, the code raises an AttributeError during the visualization phase.

This occurs because jacobian_plot receives a FunctionMetadata object as the kernel parameter, but attempts to call .get() on it.

---

Steps to Reproduce

import warp as wp
import warp.autograd

wp.init()

@wp.kernel
def mul(x: wp.array(dtype=float), s: wp.array(dtype=float), out: wp.array(dtype=float)):
    i = wp.tid()
    out[i] = x[i] * s[0]

# Wrap launches inside a Python function
def pipeline(a):
    out = wp.zeros(1, dtype=float, requires_grad=True)
    x = wp.array([3.0], dtype=float)
    wp.launch(mul, 1, inputs=[x, a], outputs=[out])
    return out

a = wp.array([2.0], dtype=float, requires_grad=True)

# Run autograd jacobian with plot_jacobians=True
jacs = wp.autograd.jacobian(
    pipeline,
    inputs=[a],
    plot_jacobians=True
)

---

Actual Behavior

Traceback (most recent call last):
  File "test.py", line 20, in <module>
    jacs = wp.autograd.jacobian(
           ^^^^^^^^^^^^^^^^^^^^^
  File "warp/_src/autograd.py", line 785, in jacobian
    jacobian_plot(
  File "warp/_src/autograd.py", line 541, in jacobian_plot
    key = kernel.key if isinstance(kernel, wp.Kernel) else kernel.get("key", "unknown")
                                                           ^^^^^^^^^^
AttributeError: 'FunctionMetadata' object has no attribute 'get'

---

Root Cause

In warp/_src/autograd.py inside jacobian_plot (around line 541):

if title is None:
    key = kernel.key if isinstance(kernel, wp.Kernel) else kernel.get("key", "unknown")

When plot_jacobians=True is used on a Python function, metadata (which is a FunctionMetadata instance) is passed as the kernel argument to jacobian_plot. Since FunctionMetadata is not a wp.Kernel and does not inherit from dict, calling .get() on it raises an AttributeError.

---

Suggested Fix

Modify warp/_src/autograd.py line 541 to safely retrieve the key attribute from FunctionMetadata or fall back to dictionary retrieval:

 if title is None:
-    key = kernel.key if isinstance(kernel, wp.Kernel) else kernel.get("key", "unknown")
+    if isinstance(kernel, wp.Kernel) or hasattr(kernel, "key"):
+        key = kernel.key or "unknown"
+    elif isinstance(kernel, dict):
+        key = kernel.get("key", "unknown")
+    else:
+        key = "unknown"
     title = f"{key} kernel Jacobian"

System Information

Environment
Warp version: 1.15.0 (and newer)
OS: Linux
Python version: 3.12

[shi-eric]

Thanks for reporting this issue, @mathisg-ui. We will try to get it fixed in the next release.