wp.sqrt() gradient is inf at x == 0, unlike wp.length() #2290

Bug Description

The adjoint of wp.sqrt() computes 0.5 / sqrt(x) with no guard, so its gradient is inf at x == 0. Downstream of that, the inf usually turns into NaN. For example, in a hand-written norm wp.sqrt(wp.dot(v, v)) at v == 0, the inf is multiplied by the zero gradient of dot:

import warp as wp


@wp.kernel
def norm_kernel(v: wp.array[wp.vec3], out: wp.array[float]):
    i = wp.tid()
    out[i] = wp.sqrt(wp.dot(v[i], v[i]))  # hand-written norm


v = wp.array([wp.vec3(3.0, 4.0, 0.0), wp.vec3(0.0, 0.0, 0.0)], dtype=wp.vec3, requires_grad=True)
out = wp.zeros(2, dtype=float, requires_grad=True)

with wp.Tape() as tape:
    wp.launch(norm_kernel, dim=2, inputs=[v], outputs=[out])
out.grad.fill_(1.0)
tape.backward()

print("out    =", out.numpy())
print("v.grad =", v.grad.numpy().tolist())

Actual output:

out    = [5. 0.]
v.grad = [[0.6000000238418579, 0.800000011920929, 0.0], [nan, nan, nan]]

Expected: v.grad[1] == [0, 0, 0], which is what wp.length(v) already returns. adj_length skips the gradient update when the result is at or below kEps, but adj_sqrt has no such guard. With wp.length() the gradient at zero is a valid subgradient (0); with an equivalent wp.sqrt() it is inf/NaN.

Proposed fix: guard adj_sqrt the same way as adj_length, giving a zero gradient at x == 0 (and -0), while still propagating NaN for x < 0, where the forward result is already NaN.

Application / Workflow Impact

MuJoCo Warp's differentiability work (google-deepmind/mujoco_warp#1727) hits this. Its enable_grad() currently monkey-patches wp.sqrt with a safe_sqrt that has a custom zero-at-origin gradient (adjoint.py). In the review discussion it was suggested that this be upstreamed to Warp instead (comment, replying to this comment). The tape-control half of that discussion is tracked separately in #2285.

Warp Version

1.20.0.dev0 (main @ b0ab769)

System Information

Python 3.12, Linux x86_64. Reproduced on CPU. The adjoint is shared C++ code in warp/native/builtin.h, so CUDA should behave the same, but I haven't tested it on CUDA.