SparseDiffTools - PolyesterForwardDiff Pre-allocation difficulties

General Usage
, ,
1

I’m trying to pre-allocate a cache to be used with SparseDiffTools and PolyesterForwardDiff, but I’m struggling to get a working example.

My attempt is based off of the jacobian tests:

difftype = ADTypes.AutoSparse(ADTypes.AutoPolyesterForwardDiff(; chunksize = 1))

J_sparsity = ForwardDiff.jacobian!(dx, obj, y, x0)

# row_colorvec = spd.matrix_colors(J_sparsity; partition_by_rows = true)
# col_colorvec = spd.matrix_colors(J_sparsity; partition_by_rows = false)
# sd = spd.PrecomputedJacobianColorvec(; jac_prototype = J_sparsity, row_colorvec, col_colorvec)

colors = spd.matrix_colors(J_sparsity)
sd = spd.PrecomputedJacobianColorvec(; jac_prototype = J_sparsity, colorvec=colors)

cache = spd.sparse_jacobian_cache(difftype, sd, obj, x0)

The final line there fails. The error message is:

ERROR: MethodError: no method matching findstructralnz(::Matrix{Float64})
The function `findstructralnz` exists, but no method is defined for this combination of argument types.

Closest candidates are:
  findstructralnz(::SparseArrays.SparseMatrixCSC)
   @ ArrayInterfaceSparseArraysExt ~/.julia/packages/ArrayInterface/9um5f/ext/ArrayInterfaceSparseArraysExt.jl:12
  findstructralnz(::Union{LinearAlgebra.SymTridiagonal, LinearAlgebra.Tridiagonal})
   @ ArrayInterface ~/.julia/packages/ArrayInterface/9um5f/src/ArrayInterface.jl:349
  findstructralnz(::LinearAlgebra.Bidiagonal)
   @ ArrayInterface ~/.julia/packages/ArrayInterface/9um5f/src/ArrayInterface.jl:341

I think I must be misunderstanding the arguments to PrecomputedJacobianColorvec. It seems that jac_prototype should be accepting an example Jacobian of your base type (so a Matrix{Float64} in my case) and the colors that you want. The example uses Symbolics to calculate the prototype jacobian, but I don’t think that’ll work with the iterative solvers in the objective function I want this to work with. My thought is to get a rough idea of the sparsity of my function by running ForwardDiff (or PolyesterForwardDiff) a couple of times. Right now, I know the sparsity for the problem I’m trying to get working, but I want to use the method I’m going to use.

For some background, obj is a functor with methods obj(y, x), and obj(x) which calculate my constraints and objective respectively. i.e.

function (obj::uo.ObjectiveFunction)(g, x)

    f = x[1]^2 - x[2] #Todo: I don't think this is needed.. 

    # g[1] = x[2] - 2*x[1]
    g[1] = - 2*x[1] #Removing x[2] to test sparsity
    g[2] = -x[2]
    g[3] = x[1]^2
    return f
end

function (obj::uo.ObjectiveFunction)(x)

    f = x[1]^2 - x[2]

    return f
end

I have those two functions because I’m going to use reverse mode for the gradient of the objective and then I want to use sparse PolyesterForwardDiff for the jacobian of the constraints.

2

Is there any reason why you would use SparseDiffTools.jl instead of the new sparse autodiff pipeline with DifferentiationInterface.jl? See the tutorial for an example.

1 Like
3

I’ll be honest, I hadn’t seen that yet. I’ll take a look at that real quick and let you know.

1 Like
4

Do reach out if there is anything that doesn’t work! You have more details on the options available in the advanced explanations of the docs.

1 Like
5

Did you end up figuring it out?

6

Sorry for the late reply, it was a travel day. It looks like a pretty nice package. I’ll get a chance to jump into it next week. I’ll hit you up with any questions I have.

It looks like it doesn’t use SparseDiffTools (not that I’m particularly attached to it). How does it’s sparsity approach vary from what’s there? Also, do you have any performance comparisons? I’m new to the JuliaDiff community, is there coordination between these two packages? Are they designed for certain use cases?

7

Sparse automatic differentiation (at least the mainstream approach discussed here) has 3 components: sparsity detection, coloring and differentiation – see this blog post for explanations. Within the last year, the corresponding Julia pipeline has been revamped in the following way to replace SparseDiffTools.jl:

task old package new package paper
sparsity detection Symbolics.jl SparseConnectivityTracer.jl [1] published
coloring SparseDiffTools.jl SparseMatrixColorings.jl submitted
differentiation ForwardDiff.jl or FiniteDiff.jl DifferentiationInterface.jl [2] submitted

The new pipeline is much faster for sparsity detection (see the paper or these benchmarks), uses more principled coloring algorithms (including ones optimized for Hessian matrices) and can work with any autodiff backend (like Enzyme.jl for instance). Most of the SciML ecosystem has transitioned to it, which means SparseDiffTools.jl might get deprecated sometime in the future. There are still some things worth optimizing (like optimized decompression for structured matrices), but unless you can measure a big performance gap in your case, I’d recommend switching to the more actively maintained and broadly used package.

  1. note that Symbolics.jl can still be used instead if necessary ↩︎

  2. includes every autodiff backend, among which ForwardDiff.jl and FiniteDiff.jl ↩︎

3 Likes
8

It already did awhile ago.

It should be noted that the new interfaces are slower for most real-world applications still because of using slower decompression on non SparseMatrixCSC, but it is much better for the health of the ecosystem for me to pass off that problem to someone else and I’m sure that will get fixed soon enough. Also the GPU support.

1 Like
9

Do you have actual benchmarks showing that? I would legitimately be curious. Feel free to add them to this issue for additional motivation.

10

Onboarding was pretty slick. :+1:

I have two things:

  • In the API for prepare_jacobian you note that the cache isn’t thread safe. Does that just apply to calling jacobian on multiple threads, or does that apply to threaded code like PolyesterForwardDiff? I would imagine the former.
  • Docs suggestion: I would include an example of using an in-place function, f(y, x). I didn’t expect to find it in the operators section of the documentation. Maybe at least point to it in a tutorial? I dunno. :man_shrugging:
1 Like
11

Yes, the former. Don’t use the same prep object from multiple threads.

The advanced tutorial has an example for in-place formulations (with the function f_sparse_vector!). Do you think another one is needed?

12

I missed that on my read-through, guess that speaks to my early morning documentation reading. I’m not sure. It doesn’t stick out as a in-place example, but it is obviously there. So not needed, but I’m not sure it would hurt to add in a small header, tag, or something.

I have a question about using the sparsity interface. The function is called once more than I would expect.

using DifferentiationInterface
import ForwardDiff
import SparseConnectivityTracer, SparseMatrixColorings

function sparse_obj_oop(x) 
    println("In sparse_obj_oop")
    o2 = - 2*x[1] 
    o3 = -x[2]
    o4 = x[1]^2

    return [o2, o3, o4]
end

sparse_forward_backend = AutoSparse(
    AutoForwardDiff();  # any object from ADTypes
    sparsity_detector=SparseConnectivityTracer.TracerSparsityDetector(),
    coloring_algorithm=SparseMatrixColorings.GreedyColoringAlgorithm(),
)

J4 = jacobian(sparse_obj_oop, sparse_forward_backend, x)

Obviously the function needs to be called for sparsity detection. But then I would expect the colored jacobian to only require one function call, but it’s taking another two function calls. Maybe I misunderstand how the coloring influences my number of function calls.

However this colors like I expect it to.

function sparse_obj(out, x)
    println("here")
    out[1] = - 2*x[1] 
    out[2] = -x[2]
    out[3] = x[1]^2
end

nx = length(x)
ny = 3  # Number of outputs

y5 = zeros(ny)

prepJ5 = prepare_jacobian(sparse_obj, y5, sparse_forward_backend, x)

J5 = jacobian(sparse_obj, y5, prepJ5, sparse_forward_backend, x)

This problem continues into my mixed-mode example I’m working through:

import ReverseDiff

function objcon(out, x) 

    println("In objcon")

    out[1] = x[1]^2 - x[2] 

    # g[1] = x[2] - 2*x[1]
    out[2] = - 2*x[1] #Removing x[2] to test sparsity
    out[3] = -x[2]
    out[4] = x[1]^2
end

sparse_mixed_backend = AutoSparse(
    MixedMode(AutoForwardDiff(), AutoReverseDiff()),
    sparsity_detector=SparseConnectivityTracer.TracerSparsityDetector(),
    coloring_algorithm=SparseMatrixColorings.GreedyColoringAlgorithm(),
)

nx = length(x)
ny = 4  # Number of outputs

y6 = zeros(ny)
prepJ6 = prepare_jacobian(objcon, y6, sparse_mixed_backend, x)

J6 = jacobian(objcon, y6, prepJ6, sparse_mixed_backend, x)
13

To count function calls, you should probably start with AutoForwardDiff(; chunksize=1) to disable chunking (otherwise some pushforwards will be batched together).
Also, the number of function calls during preparation is not a guarantee, so you should only count the number of calls during execution (not sure if that is what you were doing).
Finally, mixed mode uses a novel coloring algorithm which works best with a random vertex order, like so: GreedyColoringAlgorithm(RandomOrder(StableRNG(0), 0)). You might be one of the first people using mixed mode outside of its developers, so expect rough edges!
Can you give a full reproducible example with the same input size and choice of backend as you’re using?

14

Sorry, about that. I thought I had included all of the backends so it would be reporducible. I went and changed it. I was just using AutoForwardDiff() for my example problem.

I was counting instances after the prepare line.

I was planning on using the random initialization for my actual example, but does the algorithm not work without it? I was trying to put together the simplest working example for myself.

I’m happy to be a guinea pig!

15

How big is x? That’s the last missing piece for reproducibility.

16

Here’s what I get for case number 4 after preparation:

julia> x = ones(2)
2-element Vector{Float64}:
 1.0
 1.0

julia> prep_J4 = prepare_jacobian(sparse_obj_oop, sparse_forward_backend, x);
In sparse_obj_oop
In sparse_obj_oop

julia> J4 = jacobian(sparse_obj_oop, prep_J4, sparse_forward_backend, x);
In sparse_obj_oop

julia> ncolors(prep_J4)
1
17

Where do you see the two calls?

18

The main case here isn’t open :sweat_smile: . But I think I put something into the slack hole that was a 5x difference? I can’t find a version of it still on my computer right now though… but the profiling was pretty clearly the decompression difference so it’s at least fairly contained as to what the last gap is.

19

I don’t remember anything like that, but if you find it let me know.
I’ll try to take a look at structured decompression today.

20

Alright, sorry for the delay in response. Attached is the code that I actually ran. I re-labeled the cases for this conversation.

extra_functioncalls.jl (3.4 KB)

The cases are:

  1. Case 1: Out-of-place jacobian without preparation.
  2. Case 2: out-of-place jacobian with preparation
  3. case 3: in-place jacobian with preparation
  4. case 4: in-place mixed-mode jacobian with preparation

This sample problem is an archetype of the problem I’m trying to optimize (a wind turbine problem). The full jacobian has the gradient of the objective in the first row and and the jacobian of the constraints in the following rows. Cases 1-3 mimic just taking the partials of the constraints. Case 4 mimics the full problem.

Here are the function calls I get for each case:

  1. 3 calls
  2. 2 calls in prep, 1 call in jacobian calculation
  3. 1 call in prep, 1 call in jacobian calculation
  4. 1 call in prep, 3 calls in jacobian calculation

Case 2 suggests that 2 of the 3 calls in case 1 are in preparation. Based on case 2, I might have been counting my function calls a little differently before. Hopefully this is clearer. My questions at this point are:

  • Why is there a difference of function calls between case 2 and case 3? I expected 2 calls in case 2. What’s the second function call in preparation doing (in case 2)?
  • Any ideas why there might be 3 calls in the mixed-mode approach? I expected 1 call in prep (based on the behavior of case 3) and 2 calls in the jacobian call (1 for the reverse mode and 1 for the sparse forward call).