Best practices for use of structs with NonlinearSolve.jl

General Usage
1

I’m working on an economic model where I need to solve a system of equations.

Most of those equations are defined implicitly from a sort of mean-field game. Basically I construct several multidimensional arrays, construct an interpolated function from them, and then integrate that function over a discretized distribution. The solution of the system is the difference between those integrals and a vector of constants.

My code can be a lot neater and often faster (due to inplace mutation and stack allocation) when I organize all those arrays within a struct.

However, I’ve been struggling to get this to work well with automatic differentiation and NonlinearSolve.jl

The gist of the issue seems to be that NonlinearSolve sometimes passes a ForwardDiff dual type through only some of the variables to be solved for in the system.

As a result, I often need to do something like:

@with_kw struct mystruct{T1,T2,T3}
    A1::Array{T1,2} = fill(0.0, 2, 100)
    A2::Array{T2,2} = fill(0.0, 2, 100)
    A3::Array{T3,3} = fill(0.0, 2, 100)
end

And then within the function that sets up the system I have to do something like

function mysystem(x)
    S = mystruct{eltype(x[1]),eltype(x[2]),eltype(x[3])}()
    ...
end

This works okay but it’s cumbersome because I usually have many more than three arrays and because I don’t always know which element of x I should use to type each element of the struct.

Is it possible to somehow force the x vector to be all the same type? Alternatively, can I somehow promote non-dual type elements to dual and preserve type stability?

2

You could also create your struct with a single type parameter and promote the input types inside that struct’s (inner) constructor?

3

What about just choosing AutoEnzyme for your forward mode Jacobian here? It tends to be pretty fast and has less such restrictions. We don’t default to it yet because it has some edge cases (notably updating to v1.11) but for a majority of codes it’s likely already the better choice.

4

For some reason, Julia always crashes when I try using AutoEnzyme. Currently, I’m using version 1.11.0-rc1.

The IT administrators at work are supposed to be doing an update to v1.11.5 some time this week. Hopefully that will fix it.

Otherwise, I thought I read in one of your posts somewhere Chris - can’t remember where now unfortunately - that specifying a guess that is all dual type might solve this.

Does that sound familiar?

5

Can you provide a complete MWE?

6

I’ll try. The code library is pretty big so I’ll have to work on a smaller example and figure out how to get it to produce the same issues.

1 Like
7

This would be very useful for us to figure out if the issues are due to NonlinearSolve.jl, to its dependency DifferentiationInterface.jl, or to Enzyme.jl itself.

8

Enzyme is still having some issues with v1.11. Try v1.10 first.

9

I’m still working on getting v1.10 installed and testing out Enzyme with that.

With v1.11.5, when I run my code, before Julia crashes, I get a very long string of errors with the warning: “Warning: TODO forward zero-set of memorycopy used memset rather than runtime type…” and the stack trace leads to llvmrules.jl:792.

Not sure whether that’s helpful, but I thought I’d post it anyway in case it is.

10

Okay I got v1.10.9 installed and my script still crashes with essentially the same errors.

I’ll work on a MWE to see if I can make this easier to diagnose.

11

Can you confirm what version of Enzyme you’re on (and try the latest patch)?

12

Bypassing the autodiff backend discussion: Given the way you describe the issue, it seems the problem is that your inputs consist partly of dual numbers and partly of regular Floats. Wouldn’t something like

function mysystem(x,other_inputs)
  type_var = sum(x)
  x .+= zero(eltype(type_var))
  ...
end

provide a very simple (although perhaps unelegant) workaround? It should result in all x being promoted to the dual type eltype(type_var) and I guess in your application the cost of computing sum(x) or similar should not matter much compared to evaluating the function. Or am I naively overlooking anything here?

13

Bypassing the autodiff backend discussion: Given the way you describe the issue, it seems the problem is that your inputs consist partly of dual numbers and partly of regular Floats. Wouldn’t something like…

This is a great idea. The cost of computing the sum is indeed very small compared to evaluating the function. I’ll try it out and see how it works.

I think I’d still like to get Enzyme working though; presumably doing AD at the LLVM level should be significantly faster.

14

I’m using v0.13.53, which I think is the latest?

15

Just for fun try 0.13.54, which is the latest

16

Okay I’ve updated to 0.13.54. The problem still persists. I’m continuing to work on a MWE. Once that’s ready, hopefully this will be easier to diagnose.

1 Like