Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Use NonlinearSolve instead of Roots #155

Merged
merged 3 commits into from
Jan 7, 2021
Merged

Use NonlinearSolve instead of Roots #155

merged 3 commits into from
Jan 7, 2021

Conversation

devmotion
Copy link
Member

Motivated by #154, this PR replaces Roots with NonlinearSolve which is the root solving package that replaced Roots in the SciML ecosystem. NonlinearSolve states in its README

Performance is key: the current methods are made to be highly performant on scalar and statically sized small problems.

and

The current algorithms should support automatic differentiation, though improved adjoint overloads are planned to be added in the current update.

Copy link
Member

@cpfiffer cpfiffer left a comment

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Looks good to me, though maybe @torfjelde has opinions on it. I also agree with adding this instead of Roots -- seems they have a much more AD-friendly perspective.

@devmotion
Copy link
Member Author

Errors due to FluxML/Zygote.jl#873

@cpfiffer
Copy link
Member

cpfiffer commented Jan 5, 2021

Hm, seems like a fairly big issue -- my imagining is that thunks require some large-ish structural modficiations?

@devmotion
Copy link
Member Author

Actually I'm not sure, it could be that it is sufficient to just unthunk always when applying the results of the pullback (if it's not a thunk, unthunking would just be a no-op). But I'm not too familiar with the Zygote internals and hence I created the issue.

If we don't want to wait, my suggestion would be to just mark the failing tests as broken in DistributionsAD and Bijectors and open an issue, so we don't forget about it. IIRC it affects Skellam, Pareto, and NormalInverseGaussian with Zygote.

@cpfiffer
Copy link
Member

cpfiffer commented Jan 5, 2021

I think that'd be fine. IMO Zygote isn't really the priority AD backend anyway so I don't think we should wait on them.

Copy link
Member

@torfjelde torfjelde left a comment

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Looks good! And I'm fine with the error in Zygote.jl if we make a corresponding issue so as to remind ourselves that this needs to be updated when we have a fix.

@devmotion
Copy link
Member Author

OK, I will update the tests in TuringLang/DistributionsAD.jl#143 and then apply the same changes to the Bijectors tests which take more time to run.

Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment
Labels
None yet
Projects
None yet
Development

Successfully merging this pull request may close these issues.

3 participants