Add internal forward-mode rules for ranges #1655
Conversation
This is part 1 one solving EnzymeAD#274. It does the forward mode rules as those are simpler. A separate PR will do the WIP reverse mode rules as that seems to be a bit more complex. Add missing `@test` don't forget the rule
ChrisRackauckas
force-pushed
the
ranges_forward
branch
from
a32ccf9 to
9e976d4
Compare
src/internal_rules.jl
Outdated
| # operations as this is not directly differentiable | ||
|
|
||
| getval(x) = hasproperty(x, :val) ? x.val : x | ||
| function EnzymeRules.forward(func::Const{Colon}, ::Type{<:Duplicated}, start::Union{Const, Active}, step::Union{Const, Active}, stop::Union{Const, Active}) |
There was a problem hiding this comment.
Can you remove the union{const,active} restriction, make Type{<:Duplicated} more general, and in doing so add batchduplicated/duplicatednoneed?
See
Enzyme.jl/src/internal_rules.jl
Line 582 in d6a3ddc
Enzyme.jl/src/internal_rules.jl
Line 600 in d6a3ddc
|
How's that? |
src/internal_rules.jl
Outdated
| # operations as this is not directly differentiable | ||
|
|
||
| getval(x) = hasproperty(x, :val) ? x.val : x | ||
| function EnzymeRules.forward(func::Const{Colon}, RT::Type{<:Union{Const, DuplicatedNoNeed, Duplicated}}, start, step, stop) |
There was a problem hiding this comment.
this still needs to add batchduplicated/batchduplicatednoneed/etc , but otherwise almost there!
src/internal_rules.jl
Outdated
|
|
||
| getval(x) = hasproperty(x, :val) ? x.val : x | ||
| function EnzymeRules.forward(func::Const{Colon}, RT::Type{<:Union{Const, DuplicatedNoNeed, Duplicated}}, start, step, stop) | ||
| ret = func.val(getval.((start, step, stop))...) |
There was a problem hiding this comment.
Since it's three things can you just do .val for all of them [there isn't a question that they would have the property all Enzyme.Annotation values have .val\
|
Codecov ReportAttention: Patch coverage is
❗ Your organization needs to install the Codecov GitHub app to enable full functionality. Additional details and impacted files@@ Coverage Diff @@
## main #1655 +/- ##
==========================================
- Coverage 73.26% 71.63% -1.64%
==========================================
Files 40 31 -9
Lines 13442 13046 -396
==========================================
- Hits 9848 9345 -503
- Misses 3594 3701 +107 ☔ View full report in Codecov by Sentry. |
|
How's this? Is there anything like this |
|
You should be able to just call width(..). Also per above you should probably do something like ntuple(Val(width(RT))) do i There are several examples of this in this file |
This is part 1 one solving #274. It does the forward mode rules as those are simpler. A separate PR will do the WIP reverse mode rules as that seems to be a bit more complex.