Add internal forward-mode rules for ranges (#1655)

* Add internal forward-mode rules for ranges

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.

Add missing `@test`

don't forget the rule

* namespace

* Update internal_rules.jl

* Update internal_rules.jl

* Update src/internal_rules.jl

* Update internal_rules.jl

* Update internal_rules.jl

---------

Co-authored-by: William Moses <[email protected]>

Author: ChrisRackauckas

src/internal_rules.jl

@@ -816,6 +816,47 @@ function EnzymeRules.forward(func::Const{typeof(ldiv!)},
     end
 end
 
+# Ranges
+# Float64 ranges in Julia use bitwise `&` with higher precision
+# to correct for numerical error, thus we put rules over the
+# operations as this is not directly differentiable
+function EnzymeRules.forward(func::Const{Colon}, RT::Type{<:Union{Const, DuplicatedNoNeed, Duplicated}}, start::Annotation, step::Annotation, stop::Annotation)
+    ret = func.val(start.val, step.val, stop.val)
+    dstart = if start isa Const 
+        zero(eltype(ret)) 
+    elseif start isa Duplicated || start isa DuplicatedNoNeed
+        one(eltype(ret))
+    elseif start isa BatchDuplicated || start isa BatchDuplicatedNoNeed
+        ntuple(x->one(eltype(ret)), Val(width(RT)))
+    else
+        error("Annotation type $(typeof(start)) not supported for range start. Please open an issue")
+    end
+
+    dstep = if step isa Const 
+        zero(eltype(ret)) 
+    elseif step isa Duplicated || step isa DuplicatedNoNeed
+        one(eltype(ret))
+    elseif step isa BatchDuplicated || step isa BatchDuplicatedNoNeed
+        ntuple(x->one(eltype(ret)), Val(width(RT)))
+    else
+        error("Annotation type $(typeof(start)) not supported for range step. Please open an issue")
+    end
+
+    if RT <: Duplicated 
+        Duplicated(ret, range(dstart, step=dstep, length=length(ret)))
+    elseif RT <: Const
+        ret
+    elseif RT <: DuplicatedNoNeed
+        range(dstart, step=dstep, length=length(ret))
+    elseif RT <: BatchDuplicated
+        BatchDuplicated(ret, ntuple(x-> range(dstart, step=dstep, length=length(ret)), Val(width(RT))))
+    elseif RT <: BatchDuplicatedNoNeed
+        ntuple(x-> range(dstart, step=dstep, length=length(ret)), Val(width(RT)))
+    else
+        error("This should not be possible. Please report.")
+    end
+end
+
 function EnzymeRules.forward(
         Ty::Const{Type{BigFloat}},
         RT::Type{<:Union{DuplicatedNoNeed, Duplicated, BatchDuplicated, BatchDuplicatedNoNeed}};
@@ -876,4 +917,4 @@ function EnzymeRules.reverse(
         kwargs...,
     )
     return ()
-end
+end

test/internal_rules.jl

@@ -618,4 +618,27 @@ end
     @test autodiff(Enzyme.Reverse, x -> rand(MyDistribution(x)), Active, Active(1.0)) == ((1.0,),)
 end
 
+@testset "Ranges" begin
+    function f1(x)
+        ts = Array(0.0:x:3.0)
+        sum(ts)
+    end
+    function f2(x)
+        ts = Array(0.0:.25:3.0)
+        sum(ts) + x
+    end
+    function f3(x)
+        ts = Array(x:.25:3.0)
+        sum(ts)
+    end
+    function f4(x)
+        ts = Array(0.0:.25:x)
+        sum(ts)
+    end
+    @test Enzyme.autodiff(Forward, f1,  Duplicated(0.25, 1.0)) == (78,)
+    @test Enzyme.autodiff(Forward, f2,  Duplicated(0.25, 1.0)) == (1.0,)
+    @test Enzyme.autodiff(Forward, f3,  Duplicated(0.25, 1.0)) == (12,)
+    @test Enzyme.autodiff(Forward, f4,  Duplicated(3.0, 1.0)) == (0,)
+end
+
 end # InternalRules