Fix indexedarray adjoint

warp/_src/context.py

@@ -8120,7 +8120,11 @@ def pack_arg(kernel, arg_type, arg_name, value, device, adjoint=False):
             # - in forward passes, array types have to match
             # - in backward passes, indexed array gradients are regular arrays
             if adjoint:
-                array_matches = isinstance(value, warp.array)
+                # adjoint may be a regular array (gradient of the base storage) or the
+                # same array class as the forward arg (e.g. an indexedarray over base.grad)
+                array_matches = isinstance(value, warp.array) or (
+                    type(value) is warp._src.types.concrete_array_type(arg_type)
+                )
             else:
                 array_matches = type(value) is warp._src.types.concrete_array_type(arg_type)
 

warp/_src/types.py

@@ -4924,6 +4924,13 @@ def __str__(self):
     def __ctype__(self):
         return indexedarray_t(self.data, self.indices, self.shape)
 
+    # gradient of an indexed view is an indexed view into the base array's grad (same indices)
+    @property
+    def grad(self):
+        if self.data is None or self.data.grad is None:
+            return None
+        return indexedarray(self.data.grad, self.indices[: self.ndim], ndim=self.ndim)
+
     @property
     def vars(self):
         # member attributes available during code-gen (e.g.: d = arr.shape[0])

warp/native/array.h

@@ -1312,6 +1312,38 @@ adj_address(const array_t<T>& buf, int i, const array_t<T>& adj_buf, int adj_i,
     else if (buf.grad)
         adj_atomic_add(&index_grad(buf, i), adj_output);
 }
+// indexedarray adjoint: resolve the index indirection, then accumulate into the
+// adjoint indexedarray's storage (adj_buf.arr) or the base array's embedded grad
+template <typename T>
+inline CUDA_CALLABLE void
+adj_address(const indexedarray_t<T>& buf, int i, const indexedarray_t<T>& adj_buf, int adj_i, const T& adj_output)
+{
+    if (i < 0)
+        i += buf.shape[0];
+    if (buf.indices[0])
+        i = buf.indices[0][i];
+
+    if (adj_buf.arr.data)
+        adj_atomic_add(&index(adj_buf.arr, i), adj_output);
+    else if (buf.arr.grad)
+        adj_atomic_add(&index_grad(buf.arr, i), adj_output);
+}
+// indexedarray with a regular-array adjoint (as passed by the CUDA codegen): resolve the
+// index indirection, then accumulate into the base grad or the base array's embedded grad
+template <typename T>
+inline CUDA_CALLABLE void
+adj_address(const indexedarray_t<T>& buf, int i, const array_t<T>& adj_buf, int adj_i, const T& adj_output)
+{
+    if (i < 0)
+        i += buf.shape[0];
+    if (buf.indices[0])
+        i = buf.indices[0][i];
+
+    if (adj_buf.data)
+        adj_atomic_add(&index(adj_buf, i), adj_output);
+    else if (buf.arr.grad)
+        adj_atomic_add(&index_grad(buf.arr, i), adj_output);
+}
 template <typename T>
 inline CUDA_CALLABLE void
 adj_address(const array_t<T>& buf, int i, int j, const array_t<T>& adj_buf, int adj_i, int adj_j, const T& adj_output)

warp/tests/test_indexedarray.py

@@ -1246,6 +1246,43 @@ def test_indexedarray_fill_struct(test, device):
     assert_np_equal(a4.numpy(), np.zeros(a4.shape, dtype=nptype))
 
 
+@wp.kernel
+def kernel_indexedarray_grad_1d(
+    samples: wp.indexedarray(dtype=float),
+    weights: wp.array(dtype=float),
+    total: wp.array(dtype=float),
+):
+    i = wp.tid()
+    wp.atomic_add(total, 0, samples[i] * weights[i])
+
+
+def test_indexedarray_grad_1d(test, device):
+    # gradients must flow back through a differentiable indexedarray input (GH-1479):
+    # the adjoint follows the gather indirection and accumulates into the base array's grad
+    base = wp.array(np.linspace(1.0, 6.0, 6, dtype=np.float32), dtype=float, device=device, requires_grad=True)
+    weights_np = np.array([0.25, 0.5, 1.0], dtype=np.float32)
+    weights = wp.array(weights_np, dtype=float, device=device)
+    indices = wp.array([1, 3, 5], dtype=int, device=device)
+    samples = wp.indexedarray1d(base, [indices])
+    total = wp.zeros(1, dtype=float, device=device, requires_grad=True)
+
+    tape = wp.Tape()
+    with tape:
+        wp.launch(
+            kernel_indexedarray_grad_1d, dim=samples.size, inputs=[samples, weights], outputs=[total], device=device
+        )
+
+    # forward: sum of base[i] * weight over the gathered indices
+    assert_np_equal(total.numpy(), np.array([8.5], dtype=np.float32), tol=1e-6)
+
+    tape.backward(loss=total)
+
+    # d(total)/d(base[j]) is the matching weight at each gathered index, zero elsewhere
+    expected = np.zeros(6, dtype=np.float32)
+    expected[[1, 3, 5]] = weights_np
+    assert_np_equal(base.grad.numpy(), expected, tol=1e-6)
+
+
 devices = get_test_devices()
 
 
@@ -1254,6 +1291,7 @@ class TestIndexedArray(unittest.TestCase):
 
 
 add_function_test(TestIndexedArray, "test_indexedarray_1d", test_indexedarray_1d, devices=devices)
+add_function_test(TestIndexedArray, "test_indexedarray_grad_1d", test_indexedarray_grad_1d, devices=devices)
 add_function_test(TestIndexedArray, "test_indexedarray_2d", test_indexedarray_2d, devices=devices)
 add_function_test(TestIndexedArray, "test_indexedarray_3d", test_indexedarray_3d, devices=devices)
 add_function_test(TestIndexedArray, "test_indexedarray_4d", test_indexedarray_4d, devices=devices)