An unexpected variable reference changing during Julia Zygote gradient pullback

Viewed 116

I'm using Julia Flux.jl to train models. But when I customize my model, there is an issue the variable reference in gradient pullback function doesn't behave as expected. For simplicity, the issue code is simplified and shown as following

let
    x=rand(Float32, (10,1))
    
    v_layer=[
        Dense(10, 11),
        Dense(10, 12),
    ]
    
    ps=Flux.Params([
                vcat((collect.(params.(v_layer)))...)...,
                ])
    
    function model(x)
        out=x
        for (layer_idx, layer)=enumerate(v_layer)
            @show layer_idx, size(x)
            out=layer(x)
        end
        return out
    end
    gs=gradient(ps) do
        sum(model(x))
    end
end

If we run the code, the size of x in for loop will change according to the @show macro output as following

(layer_idx, size(x)) = (1, (10, 1))
(layer_idx, size(x)) = (2, (11, 1))

And we will get the following error message

DimensionMismatch("A has dimensions (12,10) but B has dimensions (11,1)")

I think the reference change of out in for loop should not change the reference of x, but that happened. Besides if we put the function model whole body inside gradient do as following

gs=gradient(ps) do
    out=x
    for (layer_idx, layer)=enumerate(v_layer)
        @show layer_idx, size(x)
        out=layer(x)
    end
    sum(out)
end

worked just fine. Why is this?

PS. A slightly more complex but more practical version code is as follows and has the same error.

let
    v_layer_A=[
        Dense(10, 11),
        Dense(11, 12),
    ]
    v_layer_B=[
        Dense(10, 11),
        Dense(10, 12),
    ]
    
    x=rand(Float32, (10,1))
    
    ps=Flux.Params([
                vcat((collect.(params.(v_layer_A)))...)...,
                vcat((collect.(params.(v_layer_B)))...)...,
                ])
    
    function model(x)
        out=x
        for (idx, (layer_A, layer_B))=enumerate(zip(v_layer_A, v_layer_B))
            @show idx, size(x)
            out_A=layer_A(out)
            out_B=layer_B(x)
            out=out_A.*out_B
        end
        return out
    end
    
    gs=gradient(ps) do
        sum(model(x))
    end
end
0 Answers
Related