Implementing a custom optimizer

[itk22]

Hello,
I have been using jaxopt in my latest project related to structural optimisation and I'm highly impressed with its capabilities. I have come across a specific optimization method called Svanberg's Method of Moving Asymptotes (MMA), which is known to perform well in structural optimisation tasks. To fully leverage the benefits of jaxopt and maintain a consistent interface in my codebase, I would like to implement a custom optimizer based on Svanberg's MMA within the jaxopt framework.

I would greatly appreciate your guidance on the following aspects:

• Recommendations on the relevant parts of the jaxopt codebase that I should study and modify to implement a custom optimizer.
• Advice on potential challenges or considerations I should keep in mind while integrating a new optimizer into jaxopt.

Thanks in advance!

[shoyer]

I think MMA would be a great addition to JAXopt.

If you do want to copy/adapt an existing implementation, consider using the implementation from nlopt, whose MIT license is compatible with the rest of JAXopt.

[itk22]

Hi @shoyer,
Thank you for your prompt and encouraging response. Your paper on neural reparametrization has greatly inspired my current project and I was very happy to see you respond. I'll certainly look into the nlopt implementation for MMA as a starting point.

In the meantime, could you provide me with some pointers on key aspects of the JAXopt codebase that I should focus on for implementing a custom solver? Any insight into potential challenges or best practices during the integration process would be much appreciated as well.

[shoyer]

In the meantime, could you provide me with some pointers on key aspects of the JAXopt codebase that I should focus on for implementing a custom solver? Any insight into potential challenges or best practices during the integration process would be much appreciated as well.

I haven't written much of the JAXopt code so I don't have any particular guidance.

@mblondel may have thoughts to share here

[mblondel]

Hi @acse-itk22. We don't have documentation yet on writing a custom JAXopt-compatible solver. The best would be to mimic what we did in the sources. We'll try to add some documentation soon.

[itk22]

Hi @mblondel, thank you for the response. I will have a closer look at the source code and will get back to you if I run into any issues.