Batched gradients using dr.backward()

[saeedhd96]

Hello,

I was wondering if it is possible in drjit to use ad backward mode on a vector, but avoid summation over the batch size dimension. In other words:
I realize that currently for a vector v of size B that uses a set of parameters of size n, if we use dr.backward(v), we can access the gradients in dr.grad(params) which is of size n. Is there a way to find the gradient of each row in vector v seperately which is getting gradients from dr.grad(params) of size B*n instead of n? It is sort of like the behavior of is_grads_batched in https://pytorch.org/docs/stable/generated/torch.autograd.grad.html#torch-autograd-grad.

[njroussel]

Hi @saeedhd96

We do not currently support this. To produce the same result, I believe you would need to call dr.backward(dr.gather(Float, v, index), flags=dr.ADFlag.ClearVertices) and loop over index.