Writing one component of a struct's vector or matrix field drops its gradient #2301

Bug Description

Inside a kernel, assigning to one component of a vector, matrix, quaternion, or transformation field of a local struct gives a zero gradient. The forward value is correct, and assigning the whole field is differentiable.

import warp as wp


@wp.struct
class State:
    v: wp.vec3
    t: wp.transform


@wp.kernel
def field_write(x: wp.array[float], loss: wp.array[float]):
    s = State()
    s.v = wp.vec3(x[0] * x[0], 0.0, 0.0)
    loss[0] = s.v[0]


@wp.kernel
def component_write(x: wp.array[float], loss: wp.array[float]):
    s = State()
    s.v[0] = x[0] * x[0]
    loss[0] = s.v[0]


for kernel in (field_write, component_write):
    x = wp.array([3.0], dtype=float, requires_grad=True)
    loss = wp.zeros(1, dtype=float, requires_grad=True)
    with wp.Tape() as tape:
        wp.launch(kernel, dim=1, inputs=[x, loss])
    tape.backward(loss=loss)
    print(f"{kernel.key}: loss={loss.numpy()[0]} dloss/dx={x.grad.numpy()[0]}")


@wp.kernel(module="unique")
def transform_field_write(p: wp.array[wp.vec3], out: wp.array[wp.vec3]):
    s = State()
    s.t.p = p[0]
    out[0] = s.t.p


try:
    out = wp.zeros(1, dtype=wp.vec3)
    wp.launch(transform_field_write, dim=1, inputs=[wp.ones(1, dtype=wp.vec3), out])
    print("transform_field_write:", out.numpy()[0])
except Exception as e:
    print("transform_field_write:", type(e).__name__, str(e).splitlines()[-1])

Output:

field_write: loss=9.0 dloss/dx=6.0
component_write: loss=9.0 dloss/dx=0.0
transform_field_write: WarpCodegenError ;Error, assigning transform attribute p to an array element

I expected dloss/dx=6.0 for both kernels, and s.t.p = p[0] to compile, since s is a local struct, not an array element.

The same zero gradient shows up for s.v.x = ..., s.m[1, 0] = ..., s.m[0] = row, s.q.w = ..., s.inner.v[2] = ..., s.v[1] += ..., and for a struct loaded from an array (s = arr[i]). s.v[1:] = ... fails to compile with Couldn't find function overload for 'indexref'.

The field is evaluated to a reference (&s.v), so the write is lowered to indexref + store. store is registered as non-differentiable with skip_replay, so the backward kernel neither propagates the gradient nor replays the write. Attribute stores into nested struct fields (s.inner.x = ...) avoid this in resolve_attribute_store_aggregate() by resolving the field to a member variable. That function excludes vector, matrix, quaternion, and transformation fields, and the subscript path does not use it.

I have a fix that resolves these fields to a member variable too, so component writes take the same differentiable path as writes to a local vector or matrix. I'll open a pull request that references this issue.

Application / Workflow Impact

This showed up in isolated tests that compare kernel gradients with finite differences, not in a project. It affects differentiable simulation code that keeps per-thread state in a struct with vector fields (positions, velocities, transforms) and updates single components, for example clamping one velocity axis. The gradient is silently zero.

Warp Version

main at 7c213fc (built locally, CPU only) and the 1.20.0.dev20261110 nightly. Also reproduced with 1.19.0.

System Information

macOS 26.6.2, Apple M5 Pro, Python 3.12.11, CPU. Not tried on CUDA.