ForwardColorJacCache error using partition by rows coloration

[tmigot]

I want to compute the Jacobian matrix of c! at x0 and tried to use the partition_by_rows option when computing the coloration of the Jacobian matrix as follows

c!(cx, x) = begin
    cx[1] = x[1] - 1
    cx[2] = 10 * (x[2] - x[1]^2)
    cx[3] = x[2] + 1
    cx
 end
 x0 = [-1.2; 1.0]
 nvar = 2
 ncon = 3
output = similar(x0, ncon)
sparsity_pattern = Symbolics.jacobian_sparsity(c!, output, x0)
colors = matrix_colors(sparsity_pattern, partition_by_rows = true)
ForwardColorJacCache(c!, x0, colorvec = colors, sparsity = sparsity_pattern)

but got the following error

ERROR: DimensionMismatch: row lengths must match
Stacktrace:
 [1] hcat(::BitMatrix, ::BitMatrix)
   @ Base .\bitarray.jl:1851
 [2] generate_chunked_partials(x::Vector{Float64}, colorvec::Vector{Int64}, cs::Val{2})
   @ SparseDiffTools .julia\packages\SparseDiffTools\zGdIo\src\differentiation\compute_jacobian_ad.jl:89
 [3] generate_chunked_partials(x::Vector{Float64}, colorvec::Vector{Int64}, N::Int64)
   @ SparseDiffTools .julia\packages\SparseDiffTools\zGdIo\src\differentiation\compute_jacobian_ad.jl:75
 [4] ForwardColorJacCache(f::typeof(c!), x::Vector{Float64}, _chunksize::Nothing; dx::Nothing, tag::Nothing, colorvec::Vector{Int64}, sparsity::SparseArrays.SparseMatrixCSC{Bool, Int64})
   @ SparseDiffTools .julia\packages\SparseDiffTools\zGdIo\src\differentiation\compute_jacobian_ad.jl:40
 [5] top-level scope
   @ REPL[10]:1

I also tried

ForwardColorJacCache(c!, output, colorvec = colors, sparsity = sparsity_pattern)

but then get an error when computing the jacobian with forwarddiff_color_jacobian!.

[ChrisRackauckas]

If you use partition by rows then you're setting up the coloring for reverse, not forward, differentiation.

[tmigot]

If you use partition by rows then you're setting up the coloring for reverse, not forward, differentiation.

Sorry, I forgot this. Is there a way to compute the Jacobian in reverse mode?

[ChrisRackauckas]

Yes, you just have to define the pullback vectors by the color vectors. We should make a higher level API on that like we did for forward mode, but it's just Zygote.pullback(f,u)[2](colors .== i) for each i (and don't redo the pushforward though, etc.)

[avik-pal]

We can close this now