fix: update SII syntax in CR ext also

Author: DhairyaLGandhi

ext/SciMLBaseChainRulesCoreExt.jl

@@ -3,6 +3,7 @@ module SciMLBaseChainRulesCoreExt
 using SciMLBase
 import ChainRulesCore
 import ChainRulesCore: NoTangent, @non_differentiable
+using SymbolicIndexingInterface
 
 function ChainRulesCore.rrule(
         config::ChainRulesCore.RuleConfig{
@@ -13,7 +14,7 @@ function ChainRulesCore.rrule(
         sym,
         j::Integer)
     function ODESolution_getindex_pullback(Δ)
-        i = symbolic_type(sym) != NotSymbolic() ? sym_to_index(sym, VA) : sym
+        i = symbolic_type(sym) != NotSymbolic() ? variable_index(VA, sym) : sym
         if i === nothing
             getter = getobserved(VA)
             grz = rrule_via_ad(config, getter, sym, VA.u[j], VA.prob.p, VA.t[j])[2](Δ)
@@ -66,7 +67,7 @@ end
 
 function ChainRulesCore.rrule(::typeof(getindex), VA::ODESolution, sym)
     function ODESolution_getindex_pullback(Δ)
-        i = symbolic_type(sym) != NotSymbolic() ? sym_to_index(sym, VA) : sym
+        i = symbolic_type(sym) != NotSymbolic() ? variable_index(VA, sym) : sym
         if i === nothing
             throw(error("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."))
         else

ext/SciMLBaseZygoteExt.jl

@@ -13,7 +13,7 @@ using RecursiveArrayTools
 # This method resolves the ambiguity with the pullback defined in
 # RecursiveArrayToolsZygoteExt
 # https://github.com/SciML/RecursiveArrayTools.jl/blob/d06ecb856f43bc5e37cbaf50e5f63c578bf3f1bd/ext/RecursiveArrayToolsZygoteExt.jl#L67
-@adjoint function getindex(VA::ODESolution, i::Int, j::Int)
+@adjoint function Base.getindex(VA::ODESolution, i::Int, j::Int)
     function ODESolution_getindex_pullback(Δ)
         du = [m == j ? [i == k ? Δ : zero(VA.u[1][1]) for k in 1:length(VA.u[1])] :
               zero(VA.u[1]) for m in 1:length(VA.u)]
@@ -38,7 +38,7 @@ using RecursiveArrayTools
     VA[i, j], ODESolution_getindex_pullback
 end
 
-@adjoint function getindex(VA::ODESolution, sym, j::Int)
+@adjoint function Base.getindex(VA::ODESolution, sym, j::Int)
     function ODESolution_getindex_pullback(Δ)
         i = symbolic_type(sym) != NotSymbolic() ? variable_index(VA, sym) : sym
         du, dprob = if i === nothing
@@ -92,7 +92,7 @@ end
     out, EnsembleSolution_adjoint
 end
 
-@adjoint function getindex(VA::ODESolution, i::Int)
+@adjoint function Base.getindex(VA::ODESolution, i::Int)
     function ODESolution_getindex_pullback(Δ)
         Δ′ = [(i == j ? Δ : Zygote.FillArrays.Fill(zero(eltype(x)), size(x)))
               for (x, j) in zip(VA.u, 1:length(VA))]
@@ -106,7 +106,7 @@ end
     sim.u, p̄ -> (EnsembleSolution(p̄, 0.0, true, sim.stats),)
 end
 
-@adjoint function getindex(VA::ODESolution, sym)
+@adjoint function Base.getindex(VA::ODESolution, sym)
     function ODESolution_getindex_pullback(Δ)
         i = symbolic_type(sym) != NotSymbolic() ? variable_index(VA, sym) : sym
         if i === nothing
@@ -120,7 +120,7 @@ end
     VA[sym], ODESolution_getindex_pullback
 end
 
-@adjoint function getindex(VA::ODESolution{T}, sym::Union{Tuple, AbstractVector}) where T
+@adjoint function Base.getindex(VA::ODESolution{T}, sym::Union{Tuple, AbstractVector}) 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)

test/downstream/symbol_indexing.jl

@@ -102,9 +102,9 @@ gs_sym, = Zygote.gradient(sol) do sol
 end
 idx_sym = SymbolicIndexingInterface.variable_index(sys, lorenz1.x)
 true_grad_sym = zeros(length(ModelingToolkit.unknowns(sys)))
-true_grad_sym[idx_sym] .= 1.
+true_grad_sym[idx_sym] = 1.
 
-@test "Symbolic Indexing Adjoint: Symbol" all(x -> x == true_grad_sym, gs_sym)
+@test all(map(x -> x == true_grad_sym, gs_sym))
 
 gs_vec, = Zygote.gradient(sol) do sol
 	sum(sum.(sol[[lorenz1.x, lorenz2.x]]))
@@ -113,7 +113,7 @@ idx_vecsym = SymbolicIndexingInterface.variable_index.(Ref(sys), [lorenz1.x, lor
 true_grad_vecsym = zeros(length(ModelingToolkit.unknowns(sys)))
 true_grad_vecsym[idx_vecsym] .= 1.
 
-@test "Symbolic Indexing Adjoint: Vector{Symbol}" all(x -> x == true_grad_vecsym, gs_vec)
+@test all(map(x -> x == true_grad_vecsym, gs_vec))
 
 gs_tup, = Zygote.gradient(sol) do sol
 	sum(sum.(collect.(sol[(lorenz1.x, lorenz2.x)])))
@@ -122,13 +122,13 @@ idx_tupsym = SymbolicIndexingInterface.variable_index.(Ref(sys), [lorenz1.x, lor
 true_grad_tupsym = zeros(length(ModelingToolkit.unknowns(sys)))
 true_grad_tupsym[idx_tupsym] .= 1.
 
-@test "Symbolic Indexing Adjoint: Tuple{Symbol}" all(x -> x == true_grad_tupsym, gs_tup)
+@test all(x -> x == true_grad_tupsym, gs_tup)
 
 gs_ts, = Zygote.gradient(sol) do sol
         sum(sol[[lorenz1.x, lorenz2], :])
 end
 
-@test "Symbolic Indexing Adjoint: Timeseries/ Vector{Symbol}" all(x -> x == true_grad_vecsym, gs_ts)
+@test all(x -> x == true_grad_vecsym, gs_ts)
 
 @variables q(t)[1:2] = [1.0, 2.0]
 eqs = [D(q[1]) ~ 2q[1]