add AD (Enzyme) support via EnzymeExt

[kylebeggs]

Adds support to calculate the jacobian via Enzyme autodiff #35 . This is done by adding a package extension for Enzyme, EnzymeExt.

[JoshuaLampert]

This looks pretty nice already. Great to see that Enzyme works pretty much out of the box. I only have some minor comments/suggestions.

[codecov]

Codecov Report

Attention: Patch coverage is 0% with 4 lines in your changes missing coverage. Please review.

Project coverage is 94.57%. Comparing base (e4937e2) to head (d596ef7).

Files with missing lines	Patch %	Lines
src/differentiation.jl	0.00%	4 Missing ⚠️

Additional details and impacted files

@@               Coverage Diff               @@
##           diff-select     #129      +/-   ##
===============================================
- Coverage        95.87%   94.57%   -1.31%     
===============================================
  Files               17       17              
  Lines              291      295       +4     
===============================================
  Hits               279      279              
- Misses              12       16       +4     

☔ View full report in Codecov by Sentry.
📢 Have feedback on the report? Share it here.

[mikeingold]

@kylebeggs is there any chance you’re better at git rebase than I am? With the recent merged PR’s it’s made the change summary a little challenging to parse.

[kylebeggs]

I don't think anyone is actually good at git rebase haha. I'll give it a shot. Sometimes I find it easier to just make a new PR when there the current PR changes aren't too sparse like this one. I'm traveling this weekend so it'll be Monday.

[mikeingold]

I don't think anyone is actually good at git rebase haha. I'll give it a shot. Sometimes I find it easier to just make a new PR when there the current PR changes aren't too sparse like this one. I'm traveling this weekend so it'll be Monday.

It always sounds easy enough in practice but fails miserably for me. I just went through the same with #114 which became #130 lol. No worries on timeline. I think the additions here look good.

[kylebeggs]

@mikeingold I will still work on this in the sense that a few of the tests error and need to address that, but will wait until #131 is done to not further complicate the rebasing.

[mikeingold]

@mikeingold I will still work on this in the sense that a few of the tests error and need to address that, but will wait until #131 is done to not further complicate the rebasing.

Sounds good. Sorry again for the surprise complication. I’m really looking forward to v0.16 being a huge upgrade with all of this work we’ve gotten done.

[kylebeggs]

@mikeingold @JoshuaLampert Any ideas on how to go about adding the tests? Maybe wrpa the current tests in foreach(ad_backends) do backend? Is there a way to run the test set with twice with different setups? That would be the cleanest, but I can't figure that out.

[mikeingold]

tl;dr-up-front: I'm not sure what the best answer is.

I don't think we necessarily need exhaustive coverage of the entire combination space of: integrand function output types, geometry types, integration rule type and settings, differentiation method, floating point precision spec, etc. It would be nice, but could lead to a combinatorial explosion at test time.

The DifferentiationMethods are really just used to determine differential values via jacobian (and since jacobian is currently exported), so it might be enough to just write tests that check those results directly. We could also just add an extra test to each combinations.jl test set for an integral using it.

[JoshuaLampert]

Another option would be to have separate jobs for the different backends. Meshes.jl does something similar by setting an environment variable, see here and here and here. This would save us some CI time and the different jobs can fail and succeed independently of each other making it easier to investigate where a problem is if a job fails.

[kylebeggs]

@mikeingold I think that is a sufficient approach. @JoshuaLampert I initially tried to go this route, but I realized there is no good way to run the test suite locally in that manner. I guess you can set the variable and then run the macro that runs all the tests in the REPL?

[JoshuaLampert]

Honestly, I think the easiest would be to just change the analogue of this line locally from Float64 to Float32 (or in our case from FiniteDifference() to AutoEnzyme()).

[kylebeggs]

I decided to try the route of full integration tests. If we decide the additional tests will extend the total test time too much, we can just test jacobian, but I think this would be the ideal approach. Most of the tests pass, but there are 4 or 5 that error due to Enzyme such as Cylinder, so I guess I need to write custom rules for those. Seems Enzyme does not like the intersections check -> https://github.com/JuliaGeometry/Meshes.jl/blob/master/src/geometries/primitives/cylinder.jl#L97

[mikeingold]

Currently, one snag with this approach is that several of the geometries with specialization methods are edge cases that use integral methods with hard-coded Analytical differentials. Trying to integrate those geometries with any other DifferentiationMethod is guaranteed to throw an error.

This was one of the reasons I was hoping to just eliminate Analytical altogether. They're faster for a couple of specific geometry types, but at the cost of a fair bit of added complexity. There was some benefit in code/math reduction, but largely it was to eliminate these cases where diff_method isn't truly supported.

[mikeingold]

I'm planning to write up a Discussion post with my thoughts, and some historical context, on our unit test structure/organization. I suspect it might be adequate to just add a single new test to each geometry's standard test set, something like

# For most geometries, or @test_throws where it currently fails
@test integral(f, geometry, GaussLegendre(100); diff_method=AutoEnzyme())

I think the impact of this new diff_method should be apparent for any integral of each geometry type, which would make re-running all of them redundant. And, especially with other recent performance and CI runtime improvements, we currently have a little CI runtime to spare.

[kylebeggs]

Currently, one snag with this approach is that several of the geometries with specialization methods are edge cases that use integral methods with hard-coded Analytical differentials. Trying to integrate those geometries with any other DifferentiationMethod is guaranteed to throw an error.

I left those tests to use integral, not the integral_test functions I created in Setup, this is not an issue.

This was one of the reasons I was hoping to just eliminate Analytical altogether. They're faster for a couple of specific geometry types, but at the cost of a fair bit of added complexity. There was some benefit in code/math reduction, but largely it was to eliminate these cases where diff_method isn't truly supported.

Personally, I think the added complexity is worth the speed. But that is just my opinion.

I suspect it might be adequate to just add a single new test to each geometry's standard test set, something like

This makes sense. I will change to this approach.