Hi, this is a really amazing work. I am interested in the non-linear solver. Could you please share the code of the non-linear solver part? Thank you!