Hi all, I am converting a function originally witten in python to Julia. but I have some problems working with the autodiff function provided by Zygote when I try to take the gradient of the function ( (I am able to do this seamlessly with Jax and pytorch)

I understand that the lines where I fill in arrays P_log and P_norm are the problem but I am at a loss at what to do because P[kvals[i]:kvals[i + 1]-1] has a dynamic shape and sometimes produces a jagged array (i.e for each value of i, a different sized array might be produced - hence vcat is not an option).

I will appreciate some some pointers on how to tackle this problem (Still trying to understand the inner workings of Julia)

```
function compute_total_obj(P::AbstractVector,
F::AbstractArray,
Z::AbstractArray,
Zerr_Re::AbstractArray,
Zerr_Im::AbstractArray,
LB::AbstractVector,
UB::AbstractVector,
smf::AbstractVector,
func::Function,
num_params::Integer,
num_eis::Integer,
kvals::AbstractVector,
d2m::AbstractMatrix)
# The lines causing the error
P_log = zeros(num_params, num_eis)
P_norm = zeros(num_params, num_eis)
for i = 1:num_params
P_log[i, :] .= (P[kvals[i]:kvals[i + 1]-1])
P_norm[i, :] .= ((LB[kvals[i]:kvals[i + 1]-1] .+ (10 .^ P[kvals[i]:kvals[i + 1]-1])) ./ (1 .+ (10 .^ P[kvals[i]:kvals[i + 1]-1]) ./ UB[kvals[i]:kvals[i + 1]-1]))
end
smf_1 = ifelse.(isinf.(smf), 0.0, smf)
chi_smf = sum(sum(((d2m * transpose(P_log)) .* (d2m * transpose(P_log))), dims = 1) .* smf_1)
wrss_tot = compute_wrss.(eachcol(P_norm), eachcol(F), eachcol(Z), eachcol(Zerr_Re), eachcol(Zerr_Im), func)
return (sum(wrss_tot) + chi_smf)
end
# Taking the gradient
function g!(G, p)
G .= Zygote.gradient(p -> compute_total_obj(p, F, Z, Zerr_Re, Zerr_Im, lb_vec, ub_vec, smf, func, num_params, num_eis, kvals, d2m), p)[1]
end
# How kvals is formed (if any element of smf is Inf, a jagged array is produced
function get_kvals(smf::AbstractVector, num_eis::Integer)
kvals = cumsum(insert!(ifelse.(isinf.(smf), 1, num_eis), 1, 1),)
return kvals
end
# Error
# Output exceeds the size limit. Open the full output data in a text editor
# Mutating arrays is not supported -- called copyto!(SubArray{Float64, 1, Matrix{Float64}, Tuple{Int64, Base.Slice{Base.OneTo{Int64}}}, true}, ...)
# This error occurs when you ask Zygote to differentiate operations that change
# the elements of arrays in place (e.g. setting values with x .= ...)
# Possible fixes:
# - avoid mutating operations (preferred)
# - or read the documentation and solutions for this error
# https://fluxml.ai/Zygote.jl/latest/limitations
# Stacktrace:
# [1] error(s::String)
# @ Base .\error.jl:35
# [2] _throw_mutation_error(f::Function, args::SubArray{Float64, 1, Matrix{Float64}, Tuple{Int64, Base.Slice{Base.OneTo{Int64}}}, true})
# @ Zygote C:\Users\richinex\.julia\packages\Zygote\dABKa\src\lib\array.jl:68
# [3] (::Zygote.var"#389#390"{SubArray{Float64, 1, Matrix{Float64}, Tuple{Int64, Base.Slice{Base.OneTo{Int64}}}, true}})(#unused#::Nothing)
# @ Zygote C:\Users\richinex\.julia\packages\Zygote\dABKa\src\lib\array.jl:83
# [4] (::Zygote.var"#2474#back#391"{Zygote.var"#389#390"{SubArray{Float64, 1, Matrix{Float64}, Tuple{Int64, Base.Slice{Base.OneTo{Int64}}}, true}}})(Δ::Nothing)
# @ Zygote C:\Users\richinex\.julia\packages\ZygoteRules\AIbCs\src\adjoint.jl:67
# [5] Pullback
# @ .\broadcast.jl:871 [inlined]
# [6] Pullback
# @ .\broadcast.jl:868 [inlined]
# [7] Pullback
# @ .\broadcast.jl:864 [inlined]
# ...
# @ Optim C:\Users\richinex\.julia\packages\Optim\Zq1jM\src\multivariate\optimize\interface.jl:150
# [22] fit_deterministic(P::Vector{Float64}, freq::Vector{Float64}, Z::Matrix{ComplexF64}, bounds::Vector{Vector{Float64}}, smf::Vector{Float64}, func::Function; weight::Type)
# @ Main d:\julia_projects\error_calc_comparison\multi_2.ipynb:43
# [23] top-level scope
# @ d:\julia_projects\error_calc_comparison\multi_2.ipynb:1
```