Different-typed forward grad (#2574)

* Different-typed forward grad

* fix

Author: wsmoses

src/sugar.jl

@@ -168,7 +168,7 @@ end
 @inline function onehot(x::NTuple{N,T}) where {T,N}
     onehot(NTuple{N,T})
 end
-@inline function onehot(x::NTuple{N,T}, start, endl) where {T,N}
+@inline function onehot(x::NTuple{N,T}, start::Int, endl::Int) where {T,N}
     ntuple(Val(endl - start + 1)) do i
         Base.@_inline_meta
         ntuple(Val(N)) do idx
@@ -182,6 +182,17 @@ end
     return (one(x),)
 end
 
+@inline function onehot(x::Tuple{Vararg{<:AbstractFloat}})
+    ntuple(Val(length(x))) do i
+        Base.@_inline_meta
+        ntuple(Val(length(x))) do idx
+            Base.@_inline_meta
+            T = typeof(x[idx])
+            return (i == idx) ? T(1) : T(0)
+        end
+    end
+end
+
 """
     gradient(::ReverseMode, f, args...)
 

test/sugar.jl

@@ -191,6 +191,14 @@ end
 
 end
 
+tupsq(t) = t[1]*t[2]
+
+@testset "Forward differing element tuple of floats" begin
+    res = Enzyme.gradient(Enzyme.Forward, tupsq, (3.0, Float32(2.0)))[1]
+    @test res[1] ≈ 2.0
+    @test res[2] ≈ 3.0
+end
+
 # these are used in gradient and jacobian tests
 struct InpStruct
     i1::Float64