def METHOD_NAME(variants_all_ad_rgb, shape, ray_o, ray_d):
    num_rays = 32
    exponent = 3.0
    kappa = 1e6
    ray = mi.Ray3f(ray_o, ray_d, 0.0, [])
    if shape == 'rectangle':
        scene = make_rectangle_mesh_scene()
    else:
        scene = make_sphere_mesh_scene()
    rng = mi.PCG32()
    params = mi.traverse(scene)
    key = 'mesh.vertex_positions'
    params.keep([key])
    init_vertex_pos = dr.unravel(mi.Point3f, params[key])
    theta = mi.Float(0.0)
    dr.enable_grad(theta)
    trans = mi.Vector3f([1.0, 0.0, 0.0])
    transform = mi.Transform4f.translate(theta * trans)
    positions_new = transform @ init_vertex_pos
    params[key] = dr.ravel(positions_new)
    params.update()
    dr.forward(theta, dr.ADFlag.ClearEdges)
    dr.set_label(theta, 'theta')
    dr.set_label(params, 'params')
    d, det = mi.ad.reparameterize_ray(
        scene=scene,
        rng=rng,
        ray=ray,
        params=params,
        active=True,
        num_rays=num_rays,
        kappa=kappa,
        exponent=exponent
    )
    assert d == ray.d
    assert det == 1.0
    dr.set_label(d, 'd')
    dr.set_label(det, 'det')
    # print(dr.graphviz_str(Float(1)))
    dr.enqueue(dr.ADMode.Forward, params)
    dr.traverse(mi.Float, dr.ADMode.Forward)
    grad_d = dr.grad(d)
    # Compute attached ray direction if the shape moves along the translation axis
    si_p = scene.ray_intersect(ray).p
    new_d = dr.normalize((si_p + trans - ray.o))
    # Gradient should match direction difference along the translation axis
    assert dr.allclose(dr.dot(new_d - ray.d, trans), dr.dot(grad_d, trans), atol=1e-2)
    # Other dimensions should be insignificant
    assert dr.all(dr.dot(grad_d, 1-trans) < 1e-4)