Skip to content

Commit

Permalink
fix: add adjoints for symbolic indexing with ::Colon
Browse files Browse the repository at this point in the history
  • Loading branch information
AayushSabharwal committed Jun 7, 2024
1 parent b1c999f commit 5112e27
Showing 1 changed file with 45 additions and 0 deletions.
45 changes: 45 additions & 0 deletions ext/SciMLBaseZygoteExt.jl
Original file line number Diff line number Diff line change
Expand Up @@ -125,6 +125,32 @@ end
VA[sym], ODESolution_getindex_pullback
end

@adjoint function Base.getindex(VA::ODESolution, sym, ::Colon)
function ODESolution_getindex_pullback(Δ)
i = symbolic_type(sym) != NotSymbolic() ? variable_index(VA, sym) : sym
if is_observed(VA, sym)
f = observed(VA, sym)
p = parameter_values(VA)
tunables, _, _ = SciMLStructures.canonicalize(SciMLStructures.Tunable(), p)
u = state_values(VA)
t = current_time(VA)
y, back = Zygote.pullback(u, tunables) do u, tunables
f.(u, Ref(tunables), t)

Check warning on line 138 in ext/SciMLBaseZygoteExt.jl

View check run for this annotation

Codecov / codecov/patch

ext/SciMLBaseZygoteExt.jl#L128-L138

Added lines #L128 - L138 were not covered by tests
end
gs = back(Δ)
(gs[1], nothing)
elseif i === nothing
throw(error("Zygote AD of purely-symbolic slicing for observed quantities is not yet supported. Work around this by using `A[sym,i]` to access each element sequentially in the function being differentiated."))

Check warning on line 143 in ext/SciMLBaseZygoteExt.jl

View check run for this annotation

Codecov / codecov/patch

ext/SciMLBaseZygoteExt.jl#L140-L143

Added lines #L140 - L143 were not covered by tests
else
@show i
Δ′ = [[i == k ? Δ[j] : zero(x[1]) for k in 1:length(x)]

Check warning on line 146 in ext/SciMLBaseZygoteExt.jl

View check run for this annotation

Codecov / codecov/patch

ext/SciMLBaseZygoteExt.jl#L145-L146

Added lines #L145 - L146 were not covered by tests
for (x, j) in zip(VA.u, 1:length(VA))]
(Δ′, nothing)

Check warning on line 148 in ext/SciMLBaseZygoteExt.jl

View check run for this annotation

Codecov / codecov/patch

ext/SciMLBaseZygoteExt.jl#L148

Added line #L148 was not covered by tests
end
end
VA[sym, :], ODESolution_getindex_pullback

Check warning on line 151 in ext/SciMLBaseZygoteExt.jl

View check run for this annotation

Codecov / codecov/patch

ext/SciMLBaseZygoteExt.jl#L151

Added line #L151 was not covered by tests
end

function obs_grads(VA, sym, obs_idx, Δ)
y, back = Zygote.pullback(VA) do sol
getindex.(Ref(sol), sym[obs_idx])
Expand Down Expand Up @@ -172,6 +198,25 @@ end
VA[sym], ODESolution_getindex_pullback
end

@adjoint function Base.getindex(

Check warning on line 201 in ext/SciMLBaseZygoteExt.jl

View check run for this annotation

Codecov / codecov/patch

ext/SciMLBaseZygoteExt.jl#L201

Added line #L201 was not covered by tests
VA::ODESolution{T}, sym::Union{Tuple, AbstractVector}, ::Colon) where {T}
function ODESolution_getindex_pullback(Δ)
sym = sym isa Tuple ? collect(sym) : sym
i = map(x -> symbolic_type(x) != NotSymbolic() ? variable_index(VA, x) : x, sym)

obs_idx = findall(s -> is_observed(VA, s), sym)
not_obs_idx = setdiff(1:length(sym), obs_idx)

gs_obs = obs_grads(VA, sym, isempty(obs_idx) ? nothing : obs_idx, Δ)
gs_not_obs = not_obs_grads(VA, sym, not_obs_idx, i, Δ)

a = Zygote.accum(gs_obs[1], gs_not_obs)

(a, nothing)
end
VA[sym], ODESolution_getindex_pullback
end

@adjoint function ODESolution{T1, T2, T3, T4, T5, T6, T7, T8, T9, T10, T11, T12
}(u,
args...) where {T1, T2, T3, T4, T5, T6, T7, T8,
Expand Down

0 comments on commit 5112e27

Please sign in to comment.