Type stability in Optimization.jl

Specific Domains Optimization (Mathematical)
1

I’ve got some code that solves several optimization problems using Optimization.jl. When checking for type stability with JET’s report_opt function, it outputs all sorts of errors around the call to OptimizationFunction. This can be duplicated with the rosenbrock problem from the tutorial; report_opt finds 10 errors with this simple test:

using Optimization, JET

rosenbrock(x, p) = (p[1] - x[1])^2 + p[2] * (x[2] - x[1]^2)^2

@report_opt OptimizationFunction(rosenbrock, Optimization.AutoForwardDiff())

I’ve only been using Julia for a couple of weeks so I’m hoping there’s a simple explanation. Thanks in advance!

2

Instantiating the OptimizationFunction is a setup step before launching the optimization. Some type instability during setup is usually not a concern as long as the optimization loop itself runs without type instabilities, so the developers of Optimization.jl apparently haven’t prioritized the effort to make this call type stable (assuming it’s even possible).

Keep in mind that what JET calls “errors” here are really just cases where the types of some variables and the corresponding methods had to be looked up at runtime instead of compile time. There’s nothing wrong/incorrect with the code.

1 Like
3

You cannot make it type stable with ForwardDiff, since it needs to do the chunk size choice.

2 Likes
4

I think it’s a better workflow to run the optimization with e.g. @btime from BenchmarkTools to see if there is a problem with excessive allocations (which is what type instability produces). If such problems appear, it’s time to investigate with @code_warntype or JET or similar things.

1 Like
5

There’s the simple explanation I was hoping for. Thanks!