diff --git a/warp/_src/context.py b/warp/_src/context.py index 479189b4a3..c8f1f10226 100644 --- a/warp/_src/context.py +++ b/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) diff --git a/warp/_src/types.py b/warp/_src/types.py index 2d4db5e76a..3d51439747 100644 --- a/warp/_src/types.py +++ b/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]) diff --git a/warp/native/array.h b/warp/native/array.h index fc373c774e..35ab455032 100644 --- a/warp/native/array.h +++ b/warp/native/array.h @@ -1312,6 +1312,38 @@ adj_address(const array_t& buf, int i, const array_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 +inline CUDA_CALLABLE void +adj_address(const indexedarray_t& buf, int i, const indexedarray_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 +inline CUDA_CALLABLE void +adj_address(const indexedarray_t& buf, int i, const array_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 inline CUDA_CALLABLE void adj_address(const array_t& buf, int i, int j, const array_t& adj_buf, int adj_i, int adj_j, const T& adj_output) diff --git a/warp/tests/test_indexedarray.py b/warp/tests/test_indexedarray.py index 7b8b6a8784..86419e2253 100644 --- a/warp/tests/test_indexedarray.py +++ b/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)