Merge pull request #81 from JuliaDiffEq/zygote

add some Zygote rules

Author: ChrisRackauckas

Project.toml

@@ -9,6 +9,7 @@ RecipesBase = "3cdcf5f2-1ef4-517c-9805-6587b60abb01"
 Requires = "ae029012-a4dd-5104-9daa-d747884805df"
 StaticArrays = "90137ffa-7385-5640-81b9-e52037218182"
 Statistics = "10745b16-79ce-11e8-11f9-7d13ad32a3b2"
+ZygoteRules = "700de1a5-db45-46bc-99cf-38207098b444"
 
 [compat]
 ArrayInterface = "1.2, 2.0"

src/RecursiveArrayTools.jl

@@ -3,7 +3,7 @@ __precompile__()
 module RecursiveArrayTools
 
   using Requires, RecipesBase, StaticArrays, Statistics,
-        ArrayInterface
+        ArrayInterface, ZygoteRules
 
   abstract type AbstractVectorOfArray{T, N, A} <: AbstractArray{T, N} end
   abstract type AbstractDiffEqArray{T, N, A} <: AbstractVectorOfArray{T, N, A} end
@@ -12,6 +12,7 @@ module RecursiveArrayTools
   include("vector_of_array.jl")
   include("array_partition.jl")
   include("init.jl")
+  include("zygote.jl")
 
   export VectorOfArray, DiffEqArray, AbstractVectorOfArray, AbstractDiffEqArray,
          vecarr_to_arr, vecarr_to_vectors, tuples

src/zygote.jl

@@ -0,0 +1,17 @@
+ZygoteRules.@adjoint function getindex(VA::AbstractVectorOfArray, i)
+  function AbstractVectorOfArray_getindex_adjoint(Δ)
+    Δ′ = Union{Nothing, eltype(VA.u)}[nothing for x in VA.u]
+    Δ′[i] = Δ
+    (Δ′,nothing)
+  end
+  VA[i],AbstractVectorOfArray_getindex_adjoint
+end
+
+ZygoteRules.@adjoint function getindex(VA::AbstractVectorOfArray, i, j...)
+  function AbstractVectorOfArray_getindex_adjoint(Δ)
+    Δ′ = zero(VA)
+    Δ′[i,j...] = Δ
+    (Δ′, map(_ -> nothing, i)...)
+  end
+  VA[i,j...],AbstractVectorOfArray_getindex_adjoint
+end