make_function returns wrong result with sparse Jacobian for very simple functions only

[gdalle]

julia> using FastDifferentiation

julia> x = make_variables(:x, 4)
4-element Vector{FastDifferentiation.Node}:
 x1
 x2
 x3
 x4

julia> y = diff(x)
3-element Vector{FastDifferentiation.Node}:
 (x2 - x1)
 (x3 - x2)
 (x4 - x3)

julia> jac = jacobian(y, x)
3×4 Matrix{FastDifferentiation.Node}:
  -1    1  0.0  0.0
 0.0   -1    1  0.0
 0.0  0.0   -1    1

julia> jac_sparse = sparse_jacobian(y, x)
3×4 SparseArrays.SparseMatrixCSC{FastDifferentiation.Node, Int64} with 6 stored entries:
 -1   1   ⋅  ⋅
  ⋅  -1   1  ⋅
  ⋅   ⋅  -1  1

julia> jac_exe = make_function(jac, x)
RuntimeGeneratedFunction(#=in FastDifferentiation=#, #=using FastDifferentiation=#, :((input_variables,)->begin
          #= /home/guillaume/.julia/packages/FastDifferentiation/7I3sT/src/CodeGeneration.jl:192 =#
          #= /home/guillaume/.julia/packages/FastDifferentiation/7I3sT/src/CodeGeneration.jl:192 =# @inbounds begin
                  #= /home/guillaume/.julia/packages/FastDifferentiation/7I3sT/src/CodeGeneration.jl:193 =#
                  begin
                      result_element_type = promote_type(Float64, eltype(input_variables))
                      result = Array{result_element_type}(undef, (3, 4))
                      result .= [-1.0 1.0 0.0 0.0; 0.0 -1.0 1.0 0.0; 0.0 0.0 -1.0 1.0]
                      return result
                  end
              end
      end))

julia> jac_sparse_exe = make_function(jac_sparse, x)
RuntimeGeneratedFunction(#=in FastDifferentiation=#, #=using FastDifferentiation=#, :((input_variables,)->begin
          #= /home/guillaume/.julia/packages/FastDifferentiation/7I3sT/src/CodeGeneration.jl:227 =#
          begin
              element_type = promote_type(Float64, eltype(input_variables))
              result = SparseMatrixCSC(3, 4, [1, 2, 4, 6, 7], [1, 1, 2, 2, 3, 3], zeros(element_type, 6))
              vals = nonzeros(result)
              vals .= [-1, 1, -1, 1, -1, 1]
              # return result missing here
          end
      end))

julia> jac_exe(rand(4))
3×4 Matrix{Float64}:
 -1.0   1.0   0.0  0.0
  0.0  -1.0   1.0  0.0
  0.0   0.0  -1.0  1.0

julia> jac_sparse_exe(rand(4))  # returns the vector of nonzeros instead of the result
6-element Vector{Float64}:
 -1.0
  1.0
 -1.0
  1.0
 -1.0
  1.0

I think this happens because of a missing return result line in the generated function. What's weird is that for a slightly more complicated function, this line does appear and the result is correct:

julia> g(x) = diff(x) .^ 2
g (generic function with 1 method)

julia> y = g(x)
3-element Vector{FastDifferentiation.Node}:
 ((x2 - x1) ^ 2)
 ((x3 - x2) ^ 2)
 ((x4 - x3) ^ 2)

julia> jac_sparse = sparse_jacobian(y, x)
3×4 SparseArrays.SparseMatrixCSC{FastDifferentiation.Node, Int64} with 6 stored entries:
 (-2 * (x2 - x1))   (2 * (x2 - x1))                 ⋅                ⋅
                ⋅  (-2 * (x3 - x2))   (2 * (x3 - x2))                ⋅
                ⋅                 ⋅  (-2 * (x4 - x3))  (2 * (x4 - x3))

julia> jac_sparse_exe = make_function(jac_sparse, x)
RuntimeGeneratedFunction(#=in FastDifferentiation=#, #=using FastDifferentiation=#, :((input_variables,)->begin
          #= /home/guillaume/.julia/packages/FastDifferentiation/7I3sT/src/CodeGeneration.jl:252 =#
          begin
              element_type = promote_type(Float64, eltype(input_variables))
              result = SparseMatrixCSC(3, 4, [1, 2, 4, 6, 7], [1, 1, 2, 2, 3, 3], zeros(element_type, 6))
              vals = nonzeros(result)
              var"##262" = input_variables[2] - input_variables[1]
              var"##261" = -2var"##262"
              vals[1] = var"##261"
              var"##263" = 2var"##262"
              vals[2] = var"##263"
              var"##265" = input_variables[3] - input_variables[2]
              var"##264" = -2var"##265"
              vals[3] = var"##264"
              var"##266" = 2var"##265"
              vals[4] = var"##266"
              var"##268" = input_variables[4] - input_variables[3]
              var"##267" = -2var"##268"
              vals[5] = var"##267"
              var"##269" = 2var"##268"
              vals[6] = var"##269"
              return result
          end
      end))

julia> jac_sparse_exe(rand(4))
3×4 SparseArrays.SparseMatrixCSC{Float64, Int64} with 6 stored entries:
 -0.943871  0.943871     ⋅           ⋅ 
   ⋅        0.0243471  -0.0243471    ⋅ 
   ⋅         ⋅          0.443617   -0.443617

[brianguenter]

Looks like it's related to one of the optimizations for constant jacobians. Should be an easy fix but may not be done till tomorrow since it's already late here.

[gdalle]

No worries at all!

[brianguenter]

Fixed the bug and am waiting for the PR CI to finish. Should be done in 15-20 minutes, then I'll make a new patch release. Let me know if this doesn't fix all the problems.

[brianguenter]

@gdalle the patch release with fix should be available soon (v 0.3.7). Let me know if this fixes the problem and I'll close the issue.

[gdalle]

Cool, thanks for the speedy response! I'll try it out once it's tagged :)