Fix numerical issue in PRB integrator

[DoeringChristian]

Description

This PR fixes a numerical issue in the PRB integrator gradients. When multiplying bsdf_val / dr.detach(bsdf_val) with J / dr.detach(J) inside the replace_grad function, the gradient will be expanded according to the product rule. Therefore, the differential of bsdf_val / dr.detach(bsdf_val) will be multiplied with J / dr.detach(J) which should be 1. However with numerical instabilities, this might not be the case. In this PR we are multiplying the values outside the replace_grad. Dr.Jit will therefore know to use 1 regardless of the actual value.

Fixes # (issue)

Testing

Checklist

• My code follows the style guidelines of this project
• My changes generate no new warnings
• My code also compiles for cuda_* and llvm_* variants. If you can't test this, please leave below
• I have commented my code
• I have made corresponding changes to the documentation
• I have added tests that prove my fix is effective or that my feature works
• I cleaned the commit history and removed any "Merge" commits
• I give permission that the Mitsuba 3 project may redistribute my contributions under the terms of its license

[lnuic]

LGTM

[lnuic]

Example of a scene that causes the issue. In this configuration, re-evaluating BSDF * cos(theta) differentiably at

mitsuba3/src/python/python/ad/integrators/prb.py

Line 290 in 8e624f2

bsdf_val = bsdf.eval(bsdf_ctx, si, wo, active_next)

returns 0.0, which later causes NaN values in the gradients.

import mitsuba as mi

mi.set_variant("cuda_ad_rgb")

integrator_prb = mi.load_dict({
    "type": "prb",
})

scene_dict = {
    'type': 'scene',
    'perpendicular_plane': {
        'type': 'obj',
        'filename': 'mitsuba3/resources/data/common/meshes/rectangle.obj',
        'to_world': mi.ScalarTransform4f.translate([0, 0, 1]),
        'flip_normals': True,
        'bsdf': {
            'type': 'conductor'
        }
    },
    'parallel_plane': {
        'type': 'obj',
        'filename': 'mitsuba3/resources/data/common/meshes/rectangle.obj',
        'to_world': mi.ScalarTransform4f.translate([0, 1, 0]).scale(2).rotate([1, 0, 0], 90),
    },
    'light': { 
        'type': 'constant'
    },
    'sensor': {
        'type': 'perspective',
        'to_world': mi.ScalarTransform4f().look_at(origin=[0, 0, 0], target=[0, 0, 1], up=[0, 1, 0]),
        'film': {
            'type': 'hdrfilm',
            'rfilter': { 'type': 'box' },
            'width': 128,
            'height': 128,
            'sample_border': True,
            'pixel_format': 'rgb',
            'component_format': 'float32',
        },
    }

}

scene = mi.load_dict(scene_dict)
scene_params = mi.traverse(scene)
dr.enable_grad(scene_params["perpendicular_plane.vertex_positions"])

img = mi.render(scene, scene_params, integrator=integrator_prb, spp=1)
dr.backward(img)

print(dr.grad(scene_params["perpendicular_plane.vertex_positions"]))