-
-
Notifications
You must be signed in to change notification settings - Fork 13
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
Add constant_function kwarg to AutoEnzyme #72
Changes from all commits
File filter
Filter by extension
Conversations
Jump to
Diff view
Diff view
There are no files selected for viewing
Original file line number | Diff line number | Diff line change |
---|---|---|
|
@@ -3,7 +3,7 @@ uuid = "47edcb42-4c32-4615-8424-f2b9edc5f35b" | |
authors = [ | ||
"Vaibhav Dixit <[email protected]>, Guillaume Dalle and contributors", | ||
] | ||
version = "1.5.4" | ||
version = "1.6.0" | ||
|
||
[deps] | ||
ChainRulesCore = "d360d2e6-b24c-11e9-a2a3-2a2ae2dbcce4" | ||
|
Original file line number | Diff line number | Diff line change |
---|---|---|
|
@@ -39,25 +39,63 @@ struct AutoDiffractor <: AbstractADType end | |
mode(::AutoDiffractor) = ForwardOrReverseMode() | ||
|
||
""" | ||
AutoEnzyme{M} | ||
AutoEnzyme{M,constant_function} | ||
|
||
Struct used to select the [Enzyme.jl](https://github.com/EnzymeAD/Enzyme.jl) backend for automatic differentiation. | ||
|
||
Defined by [ADTypes.jl](https://github.com/SciML/ADTypes.jl). | ||
|
||
# Constructors | ||
|
||
AutoEnzyme(; mode=nothing) | ||
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). | ||
|
||
# Fields | ||
|
||
- `mode::M`: can be either | ||
|
||
+ an object subtyping `EnzymeCore.Mode` (like `EnzymeCore.Forward` or `EnzymeCore.Reverse`) if a specific mode is required | ||
+ `nothing` to choose the best mode automatically | ||
|
||
# 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. Choose a reason for hiding this commentThe reason will be displayed to describe this comment to others. Learn more. 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. Choose a reason for hiding this commentThe reason will be displayed to describe this comment to others. Learn more. 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 There was a problem hiding this comment. Choose a reason for hiding this commentThe reason will be displayed to describe this comment to others. Learn more. I tried to be clearer in #73 |
||
An example of such a function is: | ||
|
||
```julia | ||
cache = [0.0] | ||
function f(x) | ||
cache[1] = x[1]^2 | ||
cache[1] + x[1] | ||
end | ||
``` | ||
|
||
In this case, the enclosed cache is a function of the differentiated input, and thus its values are non-constant with respect to the input. | ||
Thus, in order to compute the correct derivative of the output, the derivative must propagate through the `cache` value, and said `cache` must not be treated as constant. | ||
|
||
Conversely, the following function can treat `parameter` as a constant, because `parameter` is never modified based on the input `x`: | ||
|
||
```julia | ||
parameter = [0.0] | ||
function f(x) | ||
parameter[1] + x[1] | ||
end | ||
``` | ||
|
||
In this case, `constant_function = true` would allow the chosen differentiation system to perform extra memory and compute optimizations, under the assumption that `parameter` is kept constant. | ||
""" | ||
Base.@kwdef struct AutoEnzyme{M} <: AbstractADType | ||
mode::M = nothing | ||
struct AutoEnzyme{M, constant_function} <: AbstractADType | ||
mode::M | ||
end | ||
|
||
function AutoEnzyme(mode::M; constant_function::Bool = false) where {M} | ||
return AutoEnzyme{M, constant_function}(mode) | ||
end | ||
|
||
function AutoEnzyme(; mode::M = nothing, constant_function::Bool = false) where {M} | ||
return AutoEnzyme{M, constant_function}(mode) | ||
end | ||
|
||
mode(::AutoEnzyme) = ForwardOrReverseMode() # specialized in the extension | ||
|
There was a problem hiding this comment.
Choose a reason for hiding this comment
The reason will be displayed to describe this comment to others. Learn more.
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.
Choose a reason for hiding this comment
The reason will be displayed to describe this comment to others. Learn more.
Woops I switched true and false
There was a problem hiding this comment.
Choose a reason for hiding this comment
The reason will be displayed to describe this comment to others. Learn more.
See #73