Add constant_function kwarg to AutoEnzyme #72
Conversation
| AutoEnzyme(; mode=nothing, constant_function::Bool=false) | ||
|
|
||
| The `constant_function` keyword argument (and type parameter) determines whether the function object itself should be considered constant or not during differentiation with Enzyme.jl. | ||
| For simple functions, `constant_function` should usually be set to `false`, but in the case of closures or callable structs which contain differentiated data that can be treated as constant, `constant_function` should be set to `true` for increased performance (more details below). |
There was a problem hiding this comment.
This should read be usually set to true. However if the closure has differentiable data loaded and stored into, it may need to be set false
There was a problem hiding this comment.
Woops I switched true and false
|
|
||
| # Notes | ||
|
|
||
| If `constant_function = true` but the enclosed data is not truly constant, then Enzyme.jl will not compute the correct derivative values. |
There was a problem hiding this comment.
This should read “then Enzyme.jl will consider the load and store into cache as constant resulting in a derivative of 0”, this is still a correct derivative but subject to assumptions given to Enzyme
There was a problem hiding this comment.
I understand that you always want to stress that Enzyme computes the right thing, but this is a high-level docstring for beginners. I don't think we're implying that Enzyme is wrong here, just that someone might be using it wrong if they say constant_function=true
There was a problem hiding this comment.
I tried to be clearer in #73
Checklist
contributor guidelines, in particular the SciML Style Guide and
COLPRAC.
Additional context
This PR adds a keyword argument (and type parameter) to
AutoEnzyme, which specified whether the function itself should be considered constant during differentiation. See SciML/SciMLBenchmarks.jl#988 (comment) for the related discussion with @wsmoses and @ChrisRackauckas