Keep AD variables whose derivative evaluates to zero

[mohitt31]

Follow-up to my NEON PR #531.

I was looking at #517. The error turned out to have nothing to do with dr.square, and nothing to do with symbolic calls. It comes from the AD layer.

What happens

ad_var_new() skips an edge when its weight is a zero literal. If all the edges of an operation get skipped, the AD variable is thrown away:

if constexpr (N > 0) {
    if (!edge_index) {
        // All edges were pruned, don't create the node after all
        ad_free(ad_index, var);
        return result.release();
    }
}

After that the result has no AD index at all, so it looks the same as a variable where nobody called dr.enable_grad(). check_grad_enabled() then sees grad_enabled() == false and raises "the argument does not depend on the input variable(s) being differentiated", even though the program is fine and the gradient is simply zero.

The annoying part is that this depends on the values at runtime, so the same code fails for one input and works for another:

import drjit as dr
from drjit.llvm.ad import Float

dr.set_flag(dr.JitFlag.SymbolicCalls, True)
a, b = Float(1.0), Float(1.0)
dr.enable_grad(a)
loss = dr.square(a - b)   # a - b becomes a literal 0, so the weight 2*(a-b) is a zero literal
dr.backward(loss)         # RuntimeError, even though d(loss)/da really is 0

With Float(1.0), Float(3.0) the same code runs fine.

This also explains the SymbolicCalls part of the issue. That flag only decides whether the check inside check_grad_enabled() runs at all (it was added in 4fbd8d0). It does not change the AD graph. grad_enabled(loss) is False either way.

The change

Throw the variable away only if the operation had no differentiable input to begin with. I did not touch the edge pruning itself, so an operation whose derivative is zero stays in the graph and just does not propagate anything.

What I tested

• Full test suite on macOS arm64 (LLVM and Metal backends): 14084 passed, 0 failed.
• The new test fails on master and passes with this change, on LLVM, Metal and CUDA, for Float, Float16 and Float64. I checked CUDA separately on a T4, where the unpatched build fails all three drjit.cuda.ad variants.
• Normal AD code is not affected. The number of live AD variables is the same before and after (601 and 6001 for 200 and 2000 operations).
• The extra nodes only show up in the literal zero case. A loop of dr.square(a - b) with literal operands keeps 2 AD nodes per iteration instead of 0.

The CI failure at the moment is the nanobind dev4 / dev5 ABI mismatch that is also hitting other PRs. The build stops before any test runs.

Why I left it as a draft

I am not sure this is the trade off you want. Keeping these nodes costs memory in exactly the places the pruning was written for. The other option is to leave the AD layer alone and make the check in check_grad_enabled() less strict instead. I am happy to redo it that way if you prefer, or to add a changelog entry.

I mostly work on C++ and SIMD performance and I am still learning the JIT and autodiff internals, so I may be missing something here.