Gradient accumulation using chainrulecore

mary
Mon, 25 Sep 2023 06:04:40 +0000
New to Julia

Hi,
I am using chainrulecore and zygote and rrule to define my accumulation derivation formulation. the derivative of my function is

∑_(i=1)^n▒〖w^t  da/dr u〗

in order to solve this i have defined below code. Note that i saved all values of w,wg,ul,sl,vl in a matrix or vector.

function rrule(::typeof(gf_pf), pf_vec, vApp; β, η, fem_params)
    function U_Disp_pullback(dgdg)
      NO_FIELDS, dgdg * Dgfdpf(pf_vec,vApp; β, η, fem_params)
    end
    gf_pf(pf_vec,vApp; β, η, fem_params), U_Disp_pullback
end
function Dgfdpf(pf_vec,vApp; β, η, fem_params)
    pfh = FEFunction(fem_params.Pf, pf_vec)
    pth = (pf -> Threshold(pf; β, η)) ∘ pfh
    dgfdpf = zeros(sr)
    for i in 1:co
        u_vec = ul[:,i]
        sh_vec = sl[:,i]
        sh = FEFunction(fem_params.U_PF, sh_vec)
        vApp = vl[i]
        uAppp1(x) = VectorValue(0.0,0.0)
        uAppp2(x) = VectorValue(0.0,0.0)
        uAppp3(x) = VectorValue(0.0,-vApp)
        U_Disp = TrialFESpace(V0_Disp ,[uAppp1 ,uAppp2 ,uAppp3])
        uh = FEFunction(U_Disp, u_vec)
        wvec = w[:,i]
        wgvec = wg[:,i]
        wconjh = FEFunction(U_Disp, conj(wvec))
        wconjhg = FEFunction(fem_params.U_PF, conj(wgvec))
        l_temp(dp) = ∫( (-2*(DAdpf(wconjh,uh,pfh; β, η)) - DGdpf(wconjhg,sh,pfh; β, η) )*dp )fem_params.dΩ
        dgfdpf_temp = assemble_vector(l_temp, fem_params.Pf)
        dgfdpf = dgfdpf_temp .+ dgfdpf
    end
    return dgfdpf
end

but this code is not running and zygote cant do this.and there is error in it which is

BoundsError: attempt to access Tuple{Nothing, Nothing, Nothing, Vector{Float64}} at index [5]

Stacktrace:
  [1] getindex
    @ .\tuple.jl:29 [inlined]
  [2] gradindex(x::Tuple{Nothing, Nothing, Nothing, Vector{Float64}}, i::Int64)
    @ Zygote C:\Users\marya\.julia\packages\Zygote\IoW2g\src\compiler\reverse.jl:12
  [3] Pullback
    @ .\In[126]:3 [inlined]
  [4] (::typeof(∂(#gf_p#154)))(Δ::Float64)
    @ Zygote C:\Users\marya\.julia\packages\Zygote\IoW2g\src\compiler\interface2.jl:0
  [5] Pullback
    @ .\In[126]:1 [inlined]
  [6] (::typeof(∂(gf_p##kw)))(Δ::Float64)
    @ Zygote C:\Users\marya\.julia\packages\Zygote\IoW2g\src\compiler\interface2.jl:0
  [7] Pullback
    @ .\In[126]:8 [inlined]
  [8] (::typeof(∂(λ)))(Δ::Float64)
    @ Zygote C:\Users\marya\.julia\packages\Zygote\IoW2g\src\compiler\interface2.jl:0
  [9] (::Zygote.var"#68#69"{typeof(∂(λ))})(Δ::Float64)
    @ Zygote C:\Users\marya\.julia\packages\Zygote\IoW2g\src\compiler\interface.jl:41
 [10] gradient(f::Function, args::Vector{Float64})
    @ Zygote C:\Users\marya\.julia\packages\Zygote\IoW2g\src\compiler\interface.jl:76
 [11] gf_p(p0::Vector{Float64}, grad::Vector{Float64}, vApp::Float64; r::Float64, β::Int64, η::Float64, fem_params::NamedTuple{(:V0_Disp, :U_PF, :V0_PF, :Q, :P, :Qf, :Pf, :np, :Ω, :dΩ, :n_Γ_Load, :dΓ_Load), Tuple{Gridap.FESpaces.UnconstrainedFESpace{Vector{Float64}, Gridap.FESpaces.NodeToDofGlue{VectorValue{2, Int32}}}, Gridap.FESpaces.UnconstrainedFESpace{Vector{Float64}, Gridap.FESpaces.NodeToDofGlue{Int32}}, Gridap.FESpaces.UnconstrainedFESpace{Vector{Float64}, Gridap.FESpaces.NodeToDofGlue{Int32}}, Gridap.FESpaces.UnconstrainedFESpace{Vector{Float64}, Nothing}, Gridap.FESpaces.UnconstrainedFESpace{Vector{Float64}, Nothing}, Gridap.FESpaces.UnconstrainedFESpace{Vector{Float64}, Gridap.FESpaces.NodeToDofGlue{Int32}}, Gridap.FESpaces.UnconstrainedFESpace{Vector{Float64}, Gridap.FESpaces.NodeToDofGlue{Int32}}, Int64, BodyFittedTriangulation{2, 2, UnstructuredDiscreteModel{2, 2, Float64, Oriented}, UnstructuredGrid{2, 2, Float64, Oriented, Nothing}, Gridap.Arrays.IdentityVector{Int64}}, Gridap.CellData.GenericMeasure, GenericCellField{ReferenceDomain}, Gridap.CellData.GenericMeasure}})
    @ Main .\In[126]:8
 [12] top-level scope
    @ In[127]:5

could you please tell me how can i perform summation in zygote and chainrule? the for loop should be in gradient calculator function or it should be also in rrule function?

thank you

2 posts - 2 participants