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

Move Enzyme.jl into extension #71

Merged
merged 6 commits into from
Aug 27, 2024
Merged

Move Enzyme.jl into extension #71

merged 6 commits into from
Aug 27, 2024

Conversation

alyst
Copy link
Contributor

@alyst alyst commented Aug 22, 2024

Enzyme.jl is quite a heavy dependency (and comes with its own compatibility constraints), but it is not essential for factor rotations, as it is only used by ComponentLoss criteria family.

This PR does the following:

  • introduces internal autodiff_gradient!() method that is agnostic of a particular autodiff backend
  • adds FactorRotations.set_autodiff_backend() global config method (defaults to Enzyme)
  • makes Enzyme.jl a weak dependency, moves autodiff_gradient!(::AutodiffBackend{:Enzyme}, ...) implementation to EnzymeExt

So the only user-visible change after this PR would be to execute using Enzyme if autodiff is required.
It also opens the possibility to have alternative backends for autodiff in the future.

Copy link
Contributor

github-actions bot commented Aug 22, 2024

Benchmark Results

main 45a4b96... main/45a4b96937f1ee...
criterion_and_gradient!/FactorRotations.Biquartimin{FactorRotations.Oblique}() 0.576 ± 0.31 μs 0.41 ± 0.024 μs 1.41
criterion_and_gradient!/FactorRotations.CrawfordFerguson{FactorRotations.Oblique, Float64}(0.5) 0.391 ± 0.0079 μs 0.406 ± 0.0079 μs 0.963
criterion_and_gradient!/FactorRotations.Geomin{Float64}(0.01) 0.592 ± 0.025 μs 0.603 ± 0.019 μs 0.982
criterion_and_gradient!/FactorRotations.Infomax{FactorRotations.Oblique}() 1.38 ± 0.042 μs 1.37 ± 0.028 μs 1.01
criterion_and_gradient!/FactorRotations.MinimumEntropy() 0.302 ± 0.0079 μs 0.302 ± 0.0069 μs 0.997
criterion_and_gradient!/FactorRotations.MinimumEntropyRatio() 30.5 μs 15.3 μs 1.99
criterion_and_gradient!/FactorRotations.Oblimax{FactorRotations.Oblique}() 0.313 ± 0.0011 μs 0.315 ± 0.001 μs 0.994
criterion_and_gradient!/FactorRotations.Oblimin{FactorRotations.Oblique, Float64}(0.5) 0.392 ± 0.007 μs 0.407 ± 0.0071 μs 0.963
criterion_and_gradient!/FactorRotations.Oblimin{FactorRotations.Orthogonal, Float64}(0.5) 0.391 ± 0.008 μs 0.411 ± 0.009 μs 0.951
criterion_and_gradient!/FactorRotations.Quartimax() 0.218 ± 0 μs 0.216 ± 0.001 μs 1.01
criterion_and_gradient!/FactorRotations.Simplimax(5) 0.623 ± 0.024 μs 0.606 ± 0.01 μs 1.03
criterion_and_gradient!/FactorRotations.TandemCriterionI() 1.59 ± 0.039 μs 3.16 ± 0.22 μs 0.504
criterion_and_gradient!/FactorRotations.TandemCriterionII() 3.31 ± 0.2 μs 1.57 ± 0.026 μs 2.1
criterion_and_gradient!/FactorRotations.Varimax() 0.151 ± 0.0011 μs 0.154 ± 0.012 μs 0.981
time_to_load 0.993 ± 0.013 s 0.384 ± 0.0017 s 2.58

@p-gw
Copy link
Owner

p-gw commented Aug 22, 2024

Very nice change, I like it!

Just one thought: What do you think about using https://github.com/gdalle/DifferentiationInterface.jl directly? It's a smaller dependency compared to Enzyme and should give us a plug and play solution for all supported AutoDiff backends.

@alyst
Copy link
Contributor Author

alyst commented Aug 22, 2024

Just one thought: What do you think about using https://github.com/gdalle/DifferentiationInterface.jl directly?

I must admit I forgot there are packages for providing common autodiff interface, when I wrote this PR. Thank you for reminding!
I don't like reinventing the wheel, and it would have been nice to use this package.
Now, when I check the DifferentiationInterface.jl deps, it looks a little bit too much given that in most use cases FactorRotations.jl does not need autodiff:
it depends on ADTypes.jl (which depends on EnzymeCore.jl), SparseMatrixColorings.jl, and FillArrays.jl (which depends on PDMats.jl, which depends on SuiteSparse.jl).

The issue that motivated this PR was that in my project I use both Turing.jl and FactorRotations.jl, so I run into some obscure version conflicts, which also affected unrelated packages.
Given that autodiff is an actively developing ecosystem, I expect it would still be a potential issue if we plug in DifferentiationInterface.jl.

My suggestion is to keep this PR as is, and reconsider the support for DiffInterface.jl/ADTypes.jl if the need for more extensive autodiff support arises.
We just need to make sure that the current model (basically, set_autodiff_backend(backend) and autodiff_gradient!(::AutodiffBackend, ...) methods) would support it without breaking the compatibility -- and I think it would handle it.

Copy link
Owner

@p-gw p-gw left a comment

Choose a reason for hiding this comment

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

The doc guide implementing_rotation_methods is still failing because it relies on autodiff for MyQuartimax. Adding Enzyme as a dependency and call using Enzyme should be sufficient to make it work again. Maybe we can add a note that the autodiff backend has to be explicitly loaded.

@p-gw
Copy link
Owner

p-gw commented Aug 23, 2024

Just one thought: What do you think about using https://github.com/gdalle/DifferentiationInterface.jl directly?

I must admit I forgot there are packages for providing common autodiff interface, when I wrote this PR. Thank you for reminding! I don't like reinventing the wheel, and it would have been nice to use this package. Now, when I check the DifferentiationInterface.jl deps, it looks a little bit too much given that in most use cases FactorRotations.jl does not need autodiff: it depends on ADTypes.jl (which depends on EnzymeCore.jl), SparseMatrixColorings.jl, and FillArrays.jl (which depends on PDMats.jl, which depends on SuiteSparse.jl).

The issue that motivated this PR was that in my project I use both Turing.jl and FactorRotations.jl, so I run into some obscure version conflicts, which also affected unrelated packages. Given that autodiff is an actively developing ecosystem, I expect it would still be a potential issue if we plug in DifferentiationInterface.jl.

My suggestion is to keep this PR as is, and reconsider the support for DiffInterface.jl/ADTypes.jl if the need for more extensive autodiff support arises. We just need to make sure that the current model (basically, set_autodiff_backend(backend) and autodiff_gradient!(::AutodiffBackend, ...) methods) would support it without breaking the compatibility -- and I think it would handle it.

That's a fair assessment. I've added some minor change requests, otherwise this is good to go.

Thanks for the effort!

@p-gw p-gw merged commit e218516 into p-gw:main Aug 27, 2024
4 of 5 checks passed
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.

2 participants