Zygote gradient fails for Custom Layer

[mayor-slash]

Zygote gradient fails for Custom Layer

[avik-pal]

Verified locally that the suggested patch works:

function (m::CustomModel)(x, ps, st::NamedTuple)
    rezi_T = 1000/x[n_species+1]
    Ns = x[1:n_species]
    model_inp = vcat(log10.(Ns), rezi_T)
    k, st_k = m.k_model(model_inp, ps.k_model, st.k_model)
    ks = 10 .^ k
    eff_Ns =  vec(prod(Ns .^ -min.(stoich_m', 0.), dims=1))
    qk = eff_Ns .* ks
    dN = vec(sum(qk .* stoich_m, dims=1))
    uscp, st_t = m.t_model([rezi_T], ps.t_model, st.t_model)
    us = 1e8 .* uscp[1:n_species]
    cp = sum((1e4 .* uscp[n_species+1:n_species*2]) .* Ns)
    dU = sum(us .* dN)
    dT = -dU/cp
    result = vcat(dN,dT)
    return result, (k_model = st_k, t_model = st_t)
end

[avik-pal]

Try changing result = vcat(dN...,dT) to not use splatting. Something like vcat(dN, dT) should work. If you splat, Zygote doesn't preserve the Vector structure

For NeuralODEs most of the sensitivity algorithms assume a flat parameter structure (unlike the default return type which is a nested named tuple). See https://lux.csail.mit.edu/dev/tutorials/intermediate/1_NeuralODE#training, it describes some Adjoints which don't require the flat parameter structure.