[BUG] Array gradients should be zeroed in storage operations

Bug Description

When an array is written to, it either represents an intermediate state in the graph or a final state. In either case, it is not an input to the graph, and so although we may kickstart an adjoint chain by setting output grads to one, the forward tensor values for the same output array are merely functions of other inputs, and so end-to-end the output adjoint is zero.

We could consider accompanying array storage operations with gradient clearing in the backward pass. Doing so would arrest adjoints accumulating over (improper) inter- and intra- kernel overwrites.

Example:

import warp as wp

@wp.kernel
def test_kernel_1(x: wp.array(dtype=float), y: wp.array(dtype=float)):
    i = wp.tid()
    a = x[i]
    b = 2.0 * a
    y[i] = b


@wp.kernel
def test_kernel_2(x: wp.array(dtype=float), y: wp.array(dtype=float)):
    i = wp.tid()
    c = x[i]
    d = 3.0 * c
    y[i] = d


N = 4

x_a = wp.ones(N, dtype=float, requires_grad=True)
x_b = wp.ones(N, dtype=float, requires_grad=True)
y = wp.zeros_like(x_b)

with wp.Tape() as tape:
    wp.launch(test_kernel_1, dim=N, inputs=[x_a, y])
    wp.launch(test_kernel_2, dim=N, inputs=[x_b, y])

y.grad = wp.ones_like(y)

tape.backward()

print(y.numpy())
print(x_a.grad.numpy())
print(x_b.grad.numpy())
print(y.grad.numpy())

# prints:
# [3. 3. 3. 3.]
# [2. 2. 2. 2.]
# [3. 3. 3. 3.]
# [1. 1. 1. 1.]

# should print:
# [3. 3. 3. 3.]
# [0. 0. 0. 0.]
# [3. 3. 3. 3.]
# [0. 0. 0. 0.]

This would represent a departure from how Warp has handled gradients thus far. Currently we keep track of "partial" adjoints for non-leaf arrays and don't bother zeroing along the way. This works but it can bury user mistakes like the one above. Zeroing would ensure array overwrites don't propagate gradients to earlier inputs.

In some (admittedly edge) cases zeroing would yield a more accurate adjoint:

@wp.kernel()
def tile_math_matmul_kernel(
    ga: wp.array2d(dtype=wp.float16), gb: wp.array2d(dtype=wp.float32), gc: wp.array2d(dtype=wp.float64)
):
    i, j = wp.tid()
    a = wp.tile_load(ga, shape=(TILE_M, TILE_K), offset=(i * TILE_M, j * TILE_K))
    b = wp.tile_load(gb, shape=(TILE_K, TILE_N), offset=(i * TILE_K, j * TILE_N))
    c = wp.tile_load(gc, shape=(TILE_M, TILE_N), offset=(i * TILE_M, j * TILE_N))
    wp.tile_matmul(a, b, c, alpha=0.5, beta=-1.3)
    wp.tile_store(gc, c, offset=(i * TILE_M, j * TILE_N))


def test_tile_math_matmul(test, device):
    rng = np.random.default_rng(42)

    A = rng.random((TILE_M, TILE_K), dtype=np.float64).astype(np.float16)
    B = rng.random((TILE_K, TILE_N), dtype=np.float32)
    C = rng.random((TILE_M, TILE_N), dtype=np.float64)

    A_wp = wp.array(A, requires_grad=True, device=device)
    B_wp = wp.array(B, requires_grad=True, device=device)
    C_wp = wp.array(C, requires_grad=True, device=device)

    with wp.Tape() as tape:
        wp.launch_tiled(
            tile_math_matmul_kernel,
            dim=[1, 1],
            inputs=[A_wp, B_wp, C_wp],
            block_dim=TILE_DIM,
            device=device,
        )

    # verify forward pass
    assert_np_equal(C_wp.numpy(), 0.5 * A @ B - 1.3 * C, tol=1e-2)

    adj_C = np.ones_like(C)

    tape.backward(grads={C_wp: wp.array(adj_C, device=device)})

    assert_np_equal(A_wp.grad.numpy(), 0.5 * adj_C @ B.T, tol=1e-2)
    assert_np_equal(B_wp.grad.numpy(), 0.5 * A.T @ adj_C, tol=1e-2)
    assert_np_equal(C_wp.grad.numpy(), adj_C - 1.3 * adj_C, tol=1e-2)

If we were to update array storage behavior, then we could write down the proper adjoint for C in this case:

assert_np_equal(C_wp.grad.numpy(), -1.3 * adj_C, tol=1e-2)

But we could also create an additional array D to get the correct result.

Alternatively we could add a warning if WRITE -> WRITE is detected in a graph.

System Information

No response

`retain_grad` for `wp.array`

What It Does

With this change, Warp will zero output gradients during the backward pass after they've been read.
This is the mathematically correct behavior — it ensures that when multiple kernels write to the
same output array, only the last write propagates gradients. This matches PyTorch and JAX
semantics.

Setting retain_grad=True on an array preserves its gradient instead of zeroing it. This is
useful when you need to inspect the gradient of an intermediate or output array after calling
tape.backward().

Quick Example

import warp as wp

@wp.kernel
def multiply(x: wp.array(dtype=float), y: wp.array(dtype=float), scale: float):
    i = wp.tid()
    y[i] = scale * x[i]

x = wp.array([1.0, 2.0, 3.0], dtype=float, requires_grad=True)
y = wp.zeros(3, dtype=float, requires_grad=True, retain_grad=True)

tape = wp.Tape()
with tape:
    wp.launch(multiply, dim=3, inputs=[x, y, 2.0])

tape.backward(grads={y: wp.ones(3, dtype=float)})

print(x.grad)  # [2. 2. 2.] — gradient of 2*x w.r.t. x
print(y.grad)  # [1. 1. 1.] — preserved because retain_grad=True

Without retain_grad=True, y.grad would be all zeros after backward.

Why This Matters

Consider a simulation loop where two kernels write to the same buffer:

with tape:
    wp.launch(compute_physics, dim=N, inputs=[state, output])
    wp.launch(apply_correction, dim=N, inputs=[correction, output])

Previously, both kernels received gradients during backward, which is incorrect — the physics
kernel's output was overwritten by the correction kernel. Now, only apply_correction receives
gradients by default, because output.grad is zeroed after the correction kernel's adjoint
reads it.

When to Use `retain_grad`

Debugging: Inspect intermediate gradients to verify your computation graph
Visualization: Plot gradient magnitudes across intermediate arrays
Multi-objective: When you intentionally want gradients to flow through all writes

When Not to Use It

Input arrays: Gradients on input arrays (arrays that are read, not written) are always
preserved — retain_grad has no effect on them
Training loops: In most optimization workflows, you only need gradients on the inputs
(leaf arrays). The default zeroing behavior is correct.

Zero grads in adj_array_store and add `retain_grad` parameter to `wp.…

Fixed by commit 63fd480. This is a breaking change.

What changed

adj_array_store now zeros each destination element's gradient after reading it during the backward pass, matching PyTorch semantics. Gradients no longer leak backward across WRITE -> WRITE overwrites, so the two kernels in the repro above now produce the expected zero gradients on x_a and y:

[3. 3. 3. 3.]
[0. 0. 0. 0.]   # x_a.grad
[3. 3. 3. 3.]
[0. 0. 0. 0.]   # y.grad

The tile_matmul edge case from the issue also now produces the correct -1.3 * adj_C gradient for the in-place accumulator C.

Opt-out via `retain_grad`

y = wp.zeros(N, dtype=wp.vec3, requires_grad=True, retain_grad=True)

retain_grad=True skips the backward zeroing for that array. Use it when you want to inspect .grad on an intermediate or output array after tape.backward().

Caveat: retain_grad=True disables the zeroing that ensures only the final write to each element produces a gradient. If any element is written more than once in the forward pass (across kernels or within one kernel), gradients will be double-counted. Only use it on arrays where each element is written at most once.

Migration notes

Code that inspects .grad on arrays written to in the forward pass will see zeros after tape.backward() unless the array was created with retain_grad=True.
When calling tape.backward() multiple times, pass fresh incoming gradients via the grads= argument rather than relying on the previous run's values.
Backward passes of kernels with many per-element array writes (e.g. matrix component assignments) may run slightly slower because of the added zeroing.

See the differentiability guide for details.

Navigation Menu

Repository navigation

[BUG] Array gradients should be zeroed in storage operations #1062

Bug Description

System Information

`retain_grad` for `wp.array`

What It Does

Quick Example

Why This Matters

When to Use `retain_grad`

When Not to Use It

What changed

Opt-out via `retain_grad`

Migration notes

Metadata

Assignees

Labels

Type

Projects

Milestone

Relationships

Development

[BUG] Array gradients should be zeroed in storage operations #1062

Description

Bug Description

System Information

Activity

daedalus5 commented on Apr 9, 2026

retain_grad for wp.array

What It Does

Quick Example

Why This Matters

When to Use retain_grad

When Not to Use It

shi-eric commented on Apr 19, 2026

What changed

Opt-out via retain_grad

Migration notes

Metadata

Metadata

Assignees

Labels

Type

Projects

Milestone

Relationships

Development

Issue actions

`retain_grad` for `wp.array`

When to Use `retain_grad`

Opt-out via `retain_grad`