SceneParameter changes are not automatically tracked

[pixelsandpointers]

Dear all,

I am writing some convenience wrappers for Mitsuba and wanted to wrap the SceneParameters in a separate class. I have seen that when modifying parameters directly in the SceneParameters the changes are picked up and gradient tracing is enabled.

import mitsuba as mi
import drjit as dr

params = mi.traverse(scene)
vertices = dr.unravel(mi.Point3f, params["Object.vertex_positions"])
opt = mi.ad.Adam(lr=0.2)
opt['trans'] = mi.Point2f(0.2, -0.2)

def transform(params, opt):
      opt["trans"] = dr.clamp(opt["trans"], -0.5, 0.5)
      trafo = mi.Transform4f.translate([opt["trans"].x, opt["trans"].y, 0.0])
      # parameter still untraced. obj.grad_enabled_() => False
      params["Object.vertex_positions"] = dr.ravel(trafo @ vertices)
      # starts tracing parameter and sets obj.grad_enabled_() => True
      params.update()

This is the desired behavior and what I try to replicate.
Now, assume I have written the following:

class Wrapper:
    def __init__(self, scene):
        self.scene = scene
        self.params = mi.traverse(scene)

    def __setitem__(self, key, value):
        self.params[key] = value
        self.params.update()


    def __call__(self):
        vertices = dr.unravel(mi.Point3f, params["Object.vertex_positions"])
        trafo = mi.Transform4f.translate([0.2, -0.2, 0.0])
        # parameter still untraced. obj.grad_enabled_() => False
        self["Object.vertex_positions"] = dr.ravel(trafo @ vertices)
        # still not tracing obj.grad_enabled_() => False

To get the desired behavior, I explicitly have to set the gradient tracking manually i.e.,

def __setitem__(self, key, value):
    dr.enable_grad(value)
    self.params[key] = value

Is this intended behavior?

Thanks in advance for your answers :)

[njroussel]

Hi @pixelsandpointers

The piece of code that enables the gradient tracking here is the optimizer and not the scene parameters.

In your first snippet of code, opt['trans'] = mi.Point2f(0.2, -0.2) implicitly enables gradient tracking on opt['trans']. After that, there we define trafo = mi.Transform4f.translate([opt["trans"].x, opt["trans"].y, 0.0]) so trafo is also tracked as it depends on opt['trans'].
Finally, with the same logic, the params["your_key"] is also tracked.

In your second snippet of code, there is no optimizer to implicitly turn on gradient tracking, hence why you need to do it explicitly.

[pixelsandpointers]

@njroussel thanks for the clarification, makes a lot of sense now.