Gradient definitions & supertypes

[mcabbott]

This is inspired by SciML/SciMLSensitivity.jl#493 (comment) . I wrote many examples there, but the core surprise was that a VectorOfArrays{T} <: AbstractArray{T,3} has a gradient Vector{Matrix{T}} which is not a 3-array. If I understood correctly, it would desirable that the result be another VectorOfArray, i.e. something AbstractArray{T,3} which iterates like Vector{Matrix}.

One way to do that would be using ProjectTo. But in fact I believe the Vector{Matrix} is being explicitly generated by a rule defined here, so perhaps it is simpler just to alter that rule to wrap this?

RFC, WIP, needs tests... and needs thought about other cases:

julia> va = RecursiveArrayTools.VectorOfArray([rand(3,3), rand(3,3)]);

julia> gradient(va -> sum(va[1]), va)[1]
VectorOfArray{Float64,3}:
2-element Vector{AbstractMatrix{Float64}}:
 3×3 Fill{Float64}, with entries equal to 1.0
 [0.0 0.0 0.0; 0.0 0.0 0.0; 0.0 0.0 0.0]

julia> ans + va |> summary  # essential this work for a valid tangent
"VectorOfArray{Float64,3}"

julia> gradient((x,y) -> sum(x ./ y), va, rand(3))[1] |> summary  # should this also be changed?
"3×3×2 Array{Float64, 3}"

julia> gradient(va -> sum(va[:,1]), va)[1]  # this didn't work before, either. 
ERROR: BoundsError: attempt to access 3×3×2 VectorOfArray{Float64, 3, Vector{Matrix{Float64}}} at index [1:3, 1]

• The broadcast case could apply something like VectorOfArray(collect(eachslice(Δ))), to return a comparable object, i.e. one which broadcasts one way and iterates a different way. This would need ProjectTo. It would also apply to things like gradient(va -> sum(abs, [1 2; 3 4] * va), VectorOfArray([[1,2], [3,4]])).
• Is indexing va[:,1] expected to work? What should it return? The other way around should I think do this:

julia> gradient(va -> sum(va[1,:]), va)[1]
VectorOfArray{Float64,3}:
2-element Vector{Matrix{Float64}}:
 [1.0 0.0 0.0; 0.0 0.0 0.0; 0.0 0.0 0.0]
 [1.0 0.0 0.0; 0.0 0.0 0.0; 0.0 0.0 0.0]

[ChrisRackauckas]

why would it not be this one?

[ChrisRackauckas]

Suggested change

@show Δ′

[mcabbott]

Sorry, very WIP!

[ChrisRackauckas]

oh no worries, still worth the downstream tests

[ChrisRackauckas]

Thanks! Running downstream tests to check it out.

[ChrisRackauckas]

yeah that's a bug