Pushforward dvals in the ranges forward rule (#1663)

* Pushforward dvals in the ranges forward rule

I just realized before it merged that we're actually missing a part of the pushforward here. Even though the solution to the derivative is just the range itself, we forgot to multiply it by the dval to pushforward the derivative. It implicitly had it as one, calculating the derivative, instead of the full pushforward. The test should be updated to catch this.

* Test dual propagation 

Changes the inputs to be the same but have another operation

* fix and test batch rules and format

Author: ChrisRackauckas

src/internal_rules.jl

@@ -820,38 +820,46 @@ end
 # 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)
+function EnzymeRules.forward(func::Const{Colon},
+                             RT::Type{<:Union{Const,DuplicatedNoNeed,Duplicated,
+                                              BatchDuplicated,BatchDuplicatedNoNeed}},
+                             start::Annotation, step::Annotation, stop::Annotation)
     ret = func.val(start.val, step.val, stop.val)
-    dstart = if start isa Const 
-        zero(eltype(ret)) 
+    dstart = if start isa Const
+        zero(eltype(ret))
     elseif start isa Duplicated || start isa DuplicatedNoNeed
-        one(eltype(ret))
+        start.dval
     elseif start isa BatchDuplicated || start isa BatchDuplicatedNoNeed
-        ntuple(x->one(eltype(ret)), Val(width(RT)))
+        ntuple(i -> start.dval[i], 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)) 
+    dstep = if step isa Const
+        zero(eltype(ret))
     elseif step isa Duplicated || step isa DuplicatedNoNeed
-        one(eltype(ret))
+        step.dval
     elseif step isa BatchDuplicated || step isa BatchDuplicatedNoNeed
-        ntuple(x->one(eltype(ret)), Val(width(RT)))
+        ntuple(i -> step.dval[i], 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)))
+    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))
+        range(dstart; step=dstep, length=length(ret))
     elseif RT <: BatchDuplicated
-        BatchDuplicated(ret, ntuple(x-> range(dstart, step=dstep, length=length(ret)), Val(width(RT))))
+        BatchDuplicated(ret,
+                        ntuple(i -> range(dstart isa Number ? dstart : dstart[i];
+                                          step=dstep isa Number ? dstep : dstep[i],
+                                          length=length(ret)), Val(width(RT))))
     elseif RT <: BatchDuplicatedNoNeed
-        ntuple(x-> range(dstart, step=dstep, length=length(ret)), Val(width(RT)))
+        ntuple(i -> range(dstart isa Number ? dstart : dstart[i];
+                          step=dstep isa Number ? dstep : dstep[i],
+                          length=length(ret)), Val(width(RT)))
     else
         error("This should not be possible. Please report.")
     end

test/internal_rules.jl

@@ -620,25 +620,38 @@ end
 
 @testset "Ranges" begin
     function f1(x)
+        x = 25.0x
         ts = Array(0.0:x:3.0)
-        sum(ts)
+        return sum(ts)
     end
     function f2(x)
-        ts = Array(0.0:.25:3.0)
-        sum(ts) + x
+        x = 25.0x
+        ts = Array(0.0:0.25:3.0)
+        return sum(ts) + x
     end
     function f3(x)
-        ts = Array(x:.25:3.0)
-        sum(ts)
+        x = 25.0x
+        ts = Array(x:0.25:3.0)
+        return 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,)
+        x = 25.0x
+        ts = Array(0.0:0.25:x)
+        return sum(ts)
+    end
+    @test Enzyme.autodiff(Forward, f1, Duplicated(0.1, 1.0)) == (25.0,)
+    @test Enzyme.autodiff(Forward, f2, Duplicated(0.1, 1.0)) == (25.0,)
+    @test Enzyme.autodiff(Forward, f3, Duplicated(0.1, 1.0)) == (75.0,)
+    @test Enzyme.autodiff(Forward, f4, Duplicated(0.12, 1.0)) == (0,)
+
+    @test Enzyme.autodiff(Forward, f1, BatchDuplicated(0.1, (1.0, 2.0))) ==
+          ((var"1"=25.0, var"2"=50.0),)
+    @test Enzyme.autodiff(Forward, f2, BatchDuplicated(0.1, (1.0, 2.0))) ==
+          ((var"1"=25.0, var"2"=50.0),)
+    @test Enzyme.autodiff(Forward, f3, BatchDuplicated(0.1, (1.0, 2.0))) ==
+          ((var"1"=75.0, var"2"=150.0),)
+    @test Enzyme.autodiff(Forward, f4, BatchDuplicated(0.12, (1.0, 2.0))) ==
+          ((var"1"=0.0, var"2"=0.0),)
 end
 
 end # InternalRules