Perturbing the Camera results in incorrect Type Promotion in Backward Pass

[manuelbb-upb]

I was trying to differentiate with respect to the Camera position and therefore included the get_primary_rays(c::Camera) function in the loss function.
Using Zygote 0.4.15 (and Julia 1.3), this results in a MethodError in the reverse pass, because at some point a Camera is instanciated with a vfov of type Float64 instead of Float32.
I have attached a minimal working example and the Stacktrace that results when executing it.

I am fairly new to Julia, so my investigation of the Zygote internals was a fruitless toil.
However there is a dirty fix to this.
In src/gradients/zygote.jl change lines 306-308 from

    getproperty(c, :vfov), Δ -> (Camera(zero(c.lookfrom), zero(c.lookat), Δ, zero(c.focus),
                                        zero(c.fixedparams)), nothing)

to

    getproperty(c, :vfov), Δ -> (Camera(zero(c.lookfrom), zero(c.lookat), Float32.(Δ), zero(c.focus),
                                        zero(c.fixedparams)), nothing)

Attachments:
mwe_and_stacktrace.zip

[avik-pal]

@manuelbb-upb In the backward pass, you need to pass a Float32 value. So the last line should be changed to pullback_fn(1.0f0)

[manuelbb-upb]

Thanks, I thought that the type was not important at that point ( I only used pullback() for debugging anyways, it happened with Zygote's gradient too).
But now something interesting has happened: After a restart of my computer I cannot reproduce the error using the minimal example provided above, it works with pullback_fn(1) and pullback_fn(1.0f0).
However, my original script still does throw the MethodError. The complete script is

using RayTracer, Images, Zygote, Flux

cd(@__DIR__)

screen_size = ( w = 40, h = 40 )

orig_obj = load_obj( "teapot-low.obj" )

light = DistantLight(
    Vec3(1.0f0),
    100.0f0,
    Vec3( 0.0f0, 1.0f0, 0.0f0 )
)

cam = Camera(
    Vec3(1.0f0, 80.0f0, -1.0f0),
    Vec3(0.0f0),
    Vec3(0.0f0, 1.0f0, 0.0f0),
    45.0f0,
    1.0f0,
    screen_size...
)

function render(camera_obj)
    origin, direction = get_primary_rays( camera_obj );
    return raytrace( origin, direction, orig_obj, light, origin, 2);
end

#utility functions
raw_to_num_array( raw_render ) = zeroonenorm(reshape(hcat(raw_render.x, raw_render.y, raw_render.z),(screen_size..., 3)))
raw_to_img( raw_render ) = get_image(raw_render, screen_size...)
num_arr_to_img( num_arr ) = colorview( RGB, permutedims( num_arr , (3,2,1) ) )

function move_cam(camera_obj, direction_vec::Vec3{Array{Float32,1}} )
    # move camera relative to object instead of moving the move_object
    moved_cam = Camera(
        camera_obj.lookfrom - direction_vec,
        camera_obj.lookat,
        camera_obj.vfov,
        camera_obj.focus,
        camera_obj.fixedparams
    )
end

# render target image for loss calculation
target_raw = render(cam)
target_num_array = raw_to_num_array(target_raw)
target_img =num_arr_to_img(target_num_array)
display(target_img)

function loss_fn( direction::Vec3 )
    new_cam = move_cam(cam, direction)
    new_raw = render( new_cam )
    new_num_array = raw_to_num_array(new_raw)

    return sum( (target_num_array .- new_num_array ).^2 )
end

pertubation = Vec3( Float32.(10 * rand(3) + (1 .+ rand(3)).^2)... )
opt = ADAM(1)
for i in 1:50
    loss, back_fn = Zygote.pullback( pertubation ) do d
        loss_fn(d)
    end
    @show loss
    grad_vec = back_fn(1.0f0)
    update!(opt, pertubation, grad_vec[1])
end

Applying my fix from the first post makes it work.

[avik-pal]

Interesting, there is nothing incorrect in your script, there must be some type promotion happening internally. It will take me some time to figure out what is the exact issue.

The fix you propose is kind of bypassing the issue, but will result in problems if we want to use Float64.