Calling similar on ArrayPartition does not respect the dims argument

[garibarba]

Usage of similar in e.g.:
https://github.com/JuliaDiff/ForwardDiff.jl/blob/909976d719fdbd5fec91a159c6e2d808c45a770f/src/jacobian.jl#L149
is broken because it returns an ArrayPartition with the wrong number of elements.
This breaks later calls to e.g. reshape (see SciML/SciMLSensitivity.jl#404).

The returned array has the same size as the source ArrayPartition, and not the expected number of elements described by the dims argument(s).

This affects these two definitions:

RecursiveArrayTools.jl/src/array_partition.jl

Lines 22 to 23 in 734999e

	# ignore dims since array partitions are vectors
	Base.similar(A::ArrayPartition, dims::NTuple{N,Int}) where {N} = similar(A)

RecursiveArrayTools.jl/src/array_partition.jl

Lines 31 to 32 in 734999e

	# ignore dims since array partitions are vectors
	Base.similar(A::ArrayPartition, ::Type{T}, dims::NTuple{N,Int}) where {T,N} = similar(A, T)

[garibarba]

Some ideas:

1. Base.similar(A::ArrayPartition, dims::NTuple{N,Int}) where {N} = similar(A.x[1], dims), as long as the types of all in A.x are the same.
2. If prod(dims) is a multiple (by factor n) of the length of the A::ArrayPartition, return an ArrayPartition of n ArrayPartitions.

[ChrisRackauckas]

I think that can work. In the ragged case just error?

[agerlach]

This shows up in Zygote.jacobian too and causes issues w/ _gradcopy!
https://github.com/FluxML/Zygote.jl/blob/4e54eb79de85a47d64776b61a1a225ccdd133fc0/src/lib/grad.jl#L142

[garibarba]

I've tried 1. and 2. on a toy SteadyStateProblem and they seem to work fine. But my larger problem is crashing in a Zygote call so I'm trying to figure that out before doing a a pull request, although it seems unrelated on first sight.