Skip to content

Fix Zygote autodiff errors #530

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Merged
merged 1 commit into from
Aug 20, 2025
Merged

Fix Zygote autodiff errors #530

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
1 commit
Select commit
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
6 changes: 4 additions & 2 deletions src/time_evolution/mcsolve.jl
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -414,6 +414,8 @@ function mcsolve(
col_times = map(i -> _mc_get_jump_callback(sol[:, i]).affect!.col_times, eachindex(sol))
col_which = map(i -> _mc_get_jump_callback(sol[:, i]).affect!.col_which, eachindex(sol))

kwargs = NamedTuple(_sol_1.prob.kwargs) # Convert to NamedTuple for Zygote.jl compatibility
Copy link

Copilot AI Aug 20, 2025

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

The NamedTuple conversion could fail if _sol_1.prob.kwargs is already a NamedTuple or contains non-symbol keys. Consider using a safer conversion or checking the type first.

Suggested change
kwargs = NamedTuple(_sol_1.prob.kwargs) # Convert to NamedTuple for Zygote.jl compatibility
if isa(_sol_1.prob.kwargs, NamedTuple)
kwargs = _sol_1.prob.kwargs
else
# Convert to NamedTuple, ensuring keys are symbols
kwargs = NamedTuple{Tuple(Symbol.(keys(_sol_1.prob.kwargs)))}(values(_sol_1.prob.kwargs))
end

Copilot uses AI. Check for mistakes.

return TimeEvolutionMCSol(
ntraj,
ens_prob_mc.times,
Expand All @@ -424,7 +426,7 @@ function mcsolve(
col_which,
sol.converged,
_sol_1.alg,
_sol_1.prob.kwargs[:abstol],
_sol_1.prob.kwargs[:reltol],
kwargs.abstol,
kwargs.reltol,
)
end
6 changes: 4 additions & 2 deletions src/time_evolution/mesolve.jl
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -204,14 +204,16 @@ function mesolve(prob::TimeEvolutionProblem, alg::OrdinaryDiffEqAlgorithm = Tsit
ρt = map(ϕ -> QuantumObject(vec2mat(ϕ), type = Operator(), dims = prob.dimensions), sol.u)
end

kwargs = NamedTuple(sol.prob.kwargs) # Convert to NamedTuple for Zygote.jl compatibility
Copy link

Copilot AI Aug 20, 2025

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

The NamedTuple conversion could fail if sol.prob.kwargs is already a NamedTuple or contains non-symbol keys. Consider using a safer conversion or checking the type first.

Suggested change
kwargs = NamedTuple(sol.prob.kwargs) # Convert to NamedTuple for Zygote.jl compatibility
if sol.prob.kwargs isa NamedTuple
kwargs = sol.prob.kwargs
elseif all(k -> k isa Symbol, keys(sol.prob.kwargs))
kwargs = NamedTuple(sol.prob.kwargs)
else
error("sol.prob.kwargs contains non-symbol keys and cannot be safely converted to a NamedTuple.")
end

Copilot uses AI. Check for mistakes.

return TimeEvolutionSol(
prob.times,
sol.t,
ρt,
_get_expvals(sol, SaveFuncMESolve),
sol.retcode,
sol.alg,
sol.prob.kwargs[:abstol],
sol.prob.kwargs[:reltol],
kwargs.abstol,
kwargs.reltol,
)
end
6 changes: 4 additions & 2 deletions src/time_evolution/sesolve.jl
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -154,14 +154,16 @@ function sesolve(prob::TimeEvolutionProblem, alg::OrdinaryDiffEqAlgorithm = Tsit

ψt = map(ϕ -> QuantumObject(ϕ, type = Ket(), dims = prob.dimensions), sol.u)

kwargs = NamedTuple(sol.prob.kwargs) # Convert to NamedTuple for Zygote.jl compatibility
Copy link

Copilot AI Aug 20, 2025

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

The NamedTuple conversion could fail if sol.prob.kwargs is already a NamedTuple or contains non-symbol keys. Consider using a safer conversion or checking the type first.

Suggested change
kwargs = NamedTuple(sol.prob.kwargs) # Convert to NamedTuple for Zygote.jl compatibility
# Safely convert to NamedTuple for Zygote.jl compatibility
if isa(sol.prob.kwargs, NamedTuple)
kwargs = sol.prob.kwargs
elseif isa(sol.prob.kwargs, Dict)
# Only convert if all keys are symbols
if all(isa(k, Symbol) for k in keys(sol.prob.kwargs))
kwargs = NamedTuple(sol.prob.kwargs)
else
error("sol.prob.kwargs contains non-symbol keys and cannot be converted to NamedTuple")
end
else
# Fallback: try conversion, or use empty NamedTuple if not possible
try
kwargs = NamedTuple(sol.prob.kwargs)
catch
kwargs = (;)
end
end

Copilot uses AI. Check for mistakes.

return TimeEvolutionSol(
prob.times,
sol.t,
ψt,
_get_expvals(sol, SaveFuncSESolve),
sol.retcode,
sol.alg,
sol.prob.kwargs[:abstol],
sol.prob.kwargs[:reltol],
kwargs.abstol,
kwargs.reltol,
)
end
6 changes: 4 additions & 2 deletions src/time_evolution/smesolve.jl
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -426,6 +426,8 @@ function smesolve(
_m_expvals_sol_1 isa Nothing ? nothing : map(i -> _get_m_expvals(sol[:, i], SaveFuncSMESolve), eachindex(sol))
m_expvals = _m_expvals isa Nothing ? nothing : stack(_m_expvals, dims = 2) # Stack on dimension 2 to align with QuTiP

kwargs = NamedTuple(_sol_1.prob.kwargs) # Convert to NamedTuple for Zygote.jl compatibility
Copy link

Copilot AI Aug 20, 2025

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

The NamedTuple conversion could fail if _sol_1.prob.kwargs is already a NamedTuple or contains non-symbol keys. Consider using a safer conversion or checking the type first.

Suggested change
kwargs = NamedTuple(_sol_1.prob.kwargs) # Convert to NamedTuple for Zygote.jl compatibility
if _sol_1.prob.kwargs isa NamedTuple
kwargs = _sol_1.prob.kwargs
elseif all(k -> k isa Symbol, keys(_sol_1.prob.kwargs))
kwargs = NamedTuple(_sol_1.prob.kwargs)
else
kwargs = (;)
end

Copilot uses AI. Check for mistakes.

return TimeEvolutionStochasticSol(
ntraj,
ens_prob.times,
Expand All @@ -435,7 +437,7 @@ function smesolve(
m_expvals, # Measurement expectation values
sol.converged,
_sol_1.alg,
_sol_1.prob.kwargs[:abstol],
_sol_1.prob.kwargs[:reltol],
kwargs.abstol,
kwargs.reltol,
)
end
6 changes: 4 additions & 2 deletions src/time_evolution/ssesolve.jl
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -418,6 +418,8 @@ function ssesolve(
_m_expvals_sol_1 isa Nothing ? nothing : map(i -> _get_m_expvals(sol[:, i], SaveFuncSSESolve), eachindex(sol))
m_expvals = _m_expvals isa Nothing ? nothing : stack(_m_expvals, dims = 2)

kwargs = NamedTuple(_sol_1.prob.kwargs) # Convert to NamedTuple for Zygote.jl compatibility
Copy link

Copilot AI Aug 20, 2025

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

The NamedTuple conversion could fail if _sol_1.prob.kwargs is already a NamedTuple or contains non-symbol keys. Consider using a safer conversion or checking the type first.

Suggested change
kwargs = NamedTuple(_sol_1.prob.kwargs) # Convert to NamedTuple for Zygote.jl compatibility
kwargs = isa(_sol_1.prob.kwargs, NamedTuple) ? _sol_1.prob.kwargs :
NamedTuple{Tuple(k for k in keys(_sol_1.prob.kwargs) if isa(k, Symbol))}(
Dict(k => v for (k, v) in _sol_1.prob.kwargs if isa(k, Symbol))
) # Safely convert to NamedTuple for Zygote.jl compatibility

Copilot uses AI. Check for mistakes.

return TimeEvolutionStochasticSol(
ntraj,
ens_prob.times,
Expand All @@ -427,7 +429,7 @@ function ssesolve(
m_expvals, # Measurement expectation values
sol.converged,
_sol_1.alg,
_sol_1.prob.kwargs[:abstol],
_sol_1.prob.kwargs[:reltol],
kwargs.abstol,
kwargs.reltol,
)
end
Loading