adrhill · adrhill · May 24, 2024 · May 22, 2024 · May 22, 2024 · May 22, 2024
diff --git a/src/SparseConnectivityTracer.jl b/src/SparseConnectivityTracer.jl
@@ -25,6 +25,7 @@ include("overload_gradient.jl")
 include("overload_hessian.jl")
 include("overload_dual.jl")
 include("overload_all.jl")
+include("overload_ifelse.jl")
 
 include("pattern.jl")
 include("adtypes.jl")

diff --git a/src/conversion.jl b/src/conversion.jl
@@ -1,38 +1,36 @@
 #! format: off
 
 ## Type conversions (non-dual)
-for TT in (GradientTracer, ConnectivityTracer, HessianTracer)
-    Base.promote_rule(::Type{T}, ::Type{N}) where {T<:TT,N<:Real} = T
-    Base.promote_rule(::Type{N}, ::Type{T}) where {T<:TT,N<:Real} = T
+Base.promote_rule(::Type{T}, ::Type{N}) where {T<:AbstractTracer,N<:Real} = T
+Base.promote_rule(::Type{N}, ::Type{T}) where {T<:AbstractTracer,N<:Real} = T
 
-    Base.big(::Type{T})   where {T<:TT} = T
-    Base.widen(::Type{T}) where {T<:TT} = T
+Base.big(::Type{T})   where {T<:AbstractTracer} = T
+Base.widen(::Type{T}) where {T<:AbstractTracer} = T
 
-    Base.convert(::Type{T}, x::Real)   where {T<:TT} = myempty(T)
-    Base.convert(::Type{T}, t::T)        where {T<:TT} = t
-    Base.convert(::Type{<:Real}, t::T) where {T<:TT} = t
+Base.convert(::Type{T}, x::Real)   where {T<:AbstractTracer} = myempty(T)
+Base.convert(::Type{T}, t::T)      where {T<:AbstractTracer} = t
+Base.convert(::Type{<:Real}, t::T) where {T<:AbstractTracer} = t
 
-    ## Constants
-    # These are methods defined on types. Methods on variables are in operators.jl 
-    Base.zero(::Type{T})        where {T<:TT} = myempty(T)
-    Base.one(::Type{T})         where {T<:TT} = myempty(T)
-    Base.oneunit(::Type{T})     where {T<:TT} = myempty(T)
-    Base.typemin(::Type{T})     where {T<:TT} = myempty(T)
-    Base.typemax(::Type{T})     where {T<:TT} = myempty(T)
-    Base.eps(::Type{T})         where {T<:TT} = myempty(T)
-    Base.float(::Type{T})       where {T<:TT} = myempty(T)
-    Base.floatmin(::Type{T})    where {T<:TT} = myempty(T)
-    Base.floatmax(::Type{T})    where {T<:TT} = myempty(T)
-    Base.maxintfloat(::Type{T}) where {T<:TT} = myempty(T)
-
-    ## Array constructors
-    Base.similar(a::Array{T,1})                     where {T<:TT}   = zeros(T, size(a, 1))
-    Base.similar(a::Array{T,2})                     where {T<:TT}   = zeros(T, size(a, 1), size(a, 2))
-    Base.similar(a::Array{A,1}, ::Type{T})          where {T<:TT,A} = zeros(T, size(a, 1))
-    Base.similar(a::Array{A,2}, ::Type{T})          where {T<:TT,A} = zeros(T, size(a, 1), size(a, 2))
-    Base.similar(::Array{T}, m::Int)                where {T<:TT}   = zeros(T, m)
-    Base.similar(::Array{T}, dims::Dims{N})         where {T<:TT,N} = zeros(T, dims)
-end
+## Constants
+# These are methods defined on types. Methods on variables are in operators.jl 
+Base.zero(::Type{T})        where {T<:AbstractTracer} = myempty(T)
+Base.one(::Type{T})         where {T<:AbstractTracer} = myempty(T)
+Base.oneunit(::Type{T})     where {T<:AbstractTracer} = myempty(T)
+Base.typemin(::Type{T})     where {T<:AbstractTracer} = myempty(T)
+Base.typemax(::Type{T})     where {T<:AbstractTracer} = myempty(T)
+Base.eps(::Type{T})         where {T<:AbstractTracer} = myempty(T)
+Base.float(::Type{T})       where {T<:AbstractTracer} = myempty(T)
+Base.floatmin(::Type{T})    where {T<:AbstractTracer} = myempty(T)
+Base.floatmax(::Type{T})    where {T<:AbstractTracer} = myempty(T)
+Base.maxintfloat(::Type{T}) where {T<:AbstractTracer} = myempty(T)
+
+## Array constructors
+Base.similar(a::Array{T,1})             where {T<:AbstractTracer}   = zeros(T, size(a, 1))
+Base.similar(a::Array{T,2})             where {T<:AbstractTracer}   = zeros(T, size(a, 1), size(a, 2))
+Base.similar(a::Array{A,1}, ::Type{T})  where {T<:AbstractTracer,A} = zeros(T, size(a, 1))
+Base.similar(a::Array{A,2}, ::Type{T})  where {T<:AbstractTracer,A} = zeros(T, size(a, 1), size(a, 2))
+Base.similar(::Array{T}, m::Int)        where {T<:AbstractTracer}   = zeros(T, m)
+Base.similar(::Array{T}, dims::Dims{N}) where {T<:AbstractTracer,N} = zeros(T, dims)
 
 Base.similar(::Array, ::Type{ConnectivityTracer{C}}, dims::Dims{N}) where {C,N}   = zeros(ConnectivityTracer{C}, dims)
 Base.similar(::Array, ::Type{GradientTracer{G}},     dims::Dims{N}) where {G,N}   = zeros(GradientTracer{G},     dims)
@@ -56,8 +54,8 @@ Base.big(::Type{D})   where {P,T,D<:Dual{P,T}} = Dual{big(P),T}
 Base.widen(::Type{D}) where {P,T,D<:Dual{P,T}} = Dual{widen(P),T}
 
 Base.convert(::Type{D}, x::Real) where {P,T,D<:Dual{P,T}}           = Dual(x, myempty(T))
-Base.convert(::Type{D}, d::D)      where {P,T,D<:Dual{P,T}}           = d
-Base.convert(::Type{N}, d::D)      where {N<:Real,P,T,D<:Dual{P,T}} = Dual(convert(T, primal(d)), tracer(d))
+Base.convert(::Type{D}, d::D)    where {P,T,D<:Dual{P,T}}           = d
+Base.convert(::Type{N}, d::D)    where {N<:Real,P,T,D<:Dual{P,T}}   = Dual(convert(T, primal(d)), tracer(d))
 
 function Base.convert(::Type{Dual{P1,T}}, d::Dual{P2,T}) where {P1,P2,T} 
     return Dual(convert(P1, primal(d)), tracer(d))

diff --git a/src/overload_all.jl b/src/overload_all.jl
@@ -6,24 +6,48 @@ function overload_all(M)
             $(overload_hessian_1_to_1(M, op))
         end for op in nameof.(list_operators_1_to_1(Val(M)))
     ]
+    exprs_1_to_1_dual = [
+        quote
+            $(overload_connectivity_1_to_1_dual(M, op))
+            $(overload_gradient_1_to_1_dual(M, op))
+            $(overload_hessian_1_to_1_dual(M, op))
+        end for op in nameof.(list_operators_1_to_1(Val(M)))
+    ]
     exprs_2_to_1 = [
         quote
             $(overload_connectivity_2_to_1(M, op))
             $(overload_gradient_2_to_1(M, op))
             $(overload_hessian_2_to_1(M, op))
         end for op in nameof.(list_operators_2_to_1(Val(M)))
     ]
+    exprs_2_to_1_dual = [
+        quote
+            $(overload_connectivity_2_to_1_dual(M, op))
+            $(overload_gradient_2_to_1_dual(M, op))
+            $(overload_hessian_2_to_1_dual(M, op))
+        end for op in nameof.(list_operators_2_to_1(Val(M)))
+    ]
     exprs_1_to_2 = [
         quote
             $(overload_connectivity_1_to_2(M, op))
             $(overload_gradient_1_to_2(M, op))
             $(overload_hessian_1_to_2(M, op))
         end for op in nameof.(list_operators_1_to_2(Val(M)))
     ]
+    exprs_1_to_2_dual = [
+        quote
+            $(overload_connectivity_1_to_2_dual(M, op))
+            $(overload_gradient_1_to_2_dual(M, op))
+            $(overload_hessian_1_to_2_dual(M, op))
+        end for op in nameof.(list_operators_1_to_2(Val(M)))
+    ]
     return quote
         $(exprs_1_to_1...)
+        $(exprs_1_to_1_dual...)
         $(exprs_2_to_1...)
+        $(exprs_2_to_1_dual...)
         $(exprs_1_to_2...)
+        $(exprs_1_to_2_dual...)
     end
 end
 

diff --git a/src/overload_connectivity.jl b/src/overload_connectivity.jl
@@ -16,6 +16,12 @@ function overload_connectivity_1_to_1(M, op)
         function $M.$op(t::T) where {T<:$SCT.ConnectivityTracer}
             return $SCT.connectivity_tracer_1_to_1(t, $SCT.is_influence_zero_global($M.$op))
         end
+    end
+end
+
+function overload_connectivity_1_to_1_dual(M, op)
+    SCT = SparseConnectivityTracer
+    return quote
         function $M.$op(d::D) where {P,T<:$SCT.ConnectivityTracer,D<:$SCT.Dual{P,T}}
             x = $SCT.primal(d)
             p_out = $M.$op(x)
@@ -58,6 +64,24 @@ function overload_connectivity_2_to_1(M, op)
                 $SCT.is_influence_arg2_zero_global($M.$op),
             )
         end
+
+        function $M.$op(tx::$SCT.ConnectivityTracer, ::Real)
+            return $SCT.connectivity_tracer_1_to_1(
+                tx, $SCT.is_influence_arg1_zero_global($M.$op)
+            )
+        end
+
+        function $M.$op(::Real, ty::$SCT.ConnectivityTracer)
+            return $SCT.connectivity_tracer_1_to_1(
+                ty, $SCT.is_influence_arg2_zero_global($M.$op)
+            )
+        end
+    end
+end
+
+function overload_connectivity_2_to_1_dual(M, op)
+    SCT = SparseConnectivityTracer
+    return quote
         function $M.$op(dx::D, dy::D) where {P,T<:$SCT.ConnectivityTracer,D<:$SCT.Dual{P,T}}
             x = $SCT.primal(dx)
             y = $SCT.primal(dy)
@@ -71,11 +95,6 @@ function overload_connectivity_2_to_1(M, op)
             return $SCT.Dual(p_out, t_out)
         end
 
-        function $M.$op(tx::$SCT.ConnectivityTracer, ::Real)
-            return $SCT.connectivity_tracer_1_to_1(
-                tx, $SCT.is_influence_arg1_zero_global($M.$op)
-            )
-        end
         function $M.$op(
             dx::D, y::Real
         ) where {P,T<:$SCT.ConnectivityTracer,D<:$SCT.Dual{P,T}}
@@ -87,11 +106,6 @@ function overload_connectivity_2_to_1(M, op)
             return $SCT.Dual(p_out, t_out)
         end
 
-        function $M.$op(::Real, ty::$SCT.ConnectivityTracer)
-            return $SCT.connectivity_tracer_1_to_1(
-                ty, $SCT.is_influence_arg2_zero_global($M.$op)
-            )
-        end
         function $M.$op(
             x::Real, dy::D
         ) where {P,T<:$SCT.ConnectivityTracer,D<:$SCT.Dual{P,T}}
@@ -125,7 +139,12 @@ function overload_connectivity_1_to_2(M, op)
                 $SCT.is_influence_out2_zero_global($M.$op),
             )
         end
+    end
+end
 
+function overload_connectivity_1_to_2_dual(M, op)
+    SCT = SparseConnectivityTracer
+    return quote
         function $M.$op(d::D) where {P,T<:$SCT.ConnectivityTracer,D<:$SCT.Dual{P,T}}
             x = $SCT.primal(d)
             p1_out, p2_out = $M.$op(x)

diff --git a/src/overload_dual.jl b/src/overload_dual.jl
@@ -24,15 +24,4 @@ for fn in (:isequal, :isapprox, :isless, :(==), :(<), :(>), :(<=), :(>=))
     @eval Base.$fn(dx::D, dy::D) where {D<:Dual} = $fn(primal(dx), primal(dy))
     @eval Base.$fn(dx::D, y::Real) where {D<:Dual} = $fn(primal(dx), y)
     @eval Base.$fn(x::Real, dy::D) where {D<:Dual} = $fn(x, primal(dy))
-
-    # Error on non-dual tracers
-    @eval function Base.$fn(tx::T, ty::T) where {T<:AbstractTracer}
-        return throw(MissingPrimalError($fn, tx))
-    end
-    @eval function Base.$fn(tx::T, y::Real) where {T<:AbstractTracer}
-        return throw(MissingPrimalError($fn, tx))
-    end
-    @eval function Base.$fn(x::Real, ty::T) where {T<:AbstractTracer}
-        return throw(MissingPrimalError($fn, ty))
-    end
 end
diff --git a/src/overload_gradient.jl b/src/overload_gradient.jl
@@ -14,6 +14,12 @@ function overload_gradient_1_to_1(M, op)
         function $M.$op(t::$SCT.GradientTracer)
             return $SCT.gradient_tracer_1_to_1(t, $SCT.is_firstder_zero_global($M.$op))
         end
+    end
+end
+
+function overload_gradient_1_to_1_dual(M, op)
+    SCT = SparseConnectivityTracer
+    return quote
         function $M.$op(d::D) where {P,T<:$SCT.GradientTracer,D<:$SCT.Dual{P,T}}
             x = $SCT.primal(d)
             p_out = $M.$op(x)
@@ -56,6 +62,23 @@ function overload_gradient_2_to_1(M, op)
                 $SCT.is_firstder_arg2_zero_global($M.$op),
             )
         end
+
+        function $M.$op(tx::$SCT.GradientTracer, ::Real)
+            return $SCT.gradient_tracer_1_to_1(
+                tx, $SCT.is_firstder_arg1_zero_global($M.$op)
+            )
+        end
+
+        function $M.$op(::Real, ty::$SCT.GradientTracer)
+            return $SCT.gradient_tracer_1_to_1(
+                ty, $SCT.is_firstder_arg2_zero_global($M.$op)
+            )
+        end
+    end
+end
+function overload_gradient_2_to_1_dual(M, op)
+    SCT = SparseConnectivityTracer
+    return quote
         function $M.$op(dx::D, dy::D) where {P,T<:$SCT.GradientTracer,D<:$SCT.Dual{P,T}}
             x = $SCT.primal(dx)
             y = $SCT.primal(dy)
@@ -69,11 +92,6 @@ function overload_gradient_2_to_1(M, op)
             return $SCT.Dual(p_out, t_out)
         end
 
-        function $M.$op(tx::$SCT.GradientTracer, ::Real)
-            return $SCT.gradient_tracer_1_to_1(
-                tx, $SCT.is_firstder_arg1_zero_global($M.$op)
-            )
-        end
         function $M.$op(dx::D, y::Real) where {P,T<:$SCT.GradientTracer,D<:$SCT.Dual{P,T}}
             x = $SCT.primal(dx)
             p_out = $M.$op(x, y)
@@ -83,11 +101,6 @@ function overload_gradient_2_to_1(M, op)
             return $SCT.Dual(p_out, t_out)
         end
 
-        function $M.$op(::Real, ty::$SCT.GradientTracer)
-            return $SCT.gradient_tracer_1_to_1(
-                ty, $SCT.is_firstder_arg2_zero_global($M.$op)
-            )
-        end
         function $M.$op(x::Real, dy::D) where {P,T<:$SCT.GradientTracer,D<:$SCT.Dual{P,T}}
             y = $SCT.primal(dy)
             p_out = $M.$op(x, y)
@@ -119,7 +132,12 @@ function overload_gradient_1_to_2(M, op)
                 $SCT.is_firstder_out2_zero_global($M.$op),
             )
         end
+    end
+end
 
+function overload_gradient_1_to_2_dual(M, op)
+    SCT = SparseConnectivityTracer
+    return quote
         function $M.$op(d::D) where {P,T<:$SCT.GradientTracer,D<:$SCT.Dual{P,T}}
             x = $SCT.primal(d)
             p1_out, p2_out = $M.$op(x)

diff --git a/src/overload_hessian.jl b/src/overload_hessian.jl
@@ -28,6 +28,12 @@ function overload_hessian_1_to_1(M, op)
                 $SCT.is_seconder_zero_global($M.$op),
             )
         end
+    end
+end
+
+function overload_hessian_1_to_1_dual(M, op)
+    SCT = SparseConnectivityTracer
+    return quote
         function $M.$op(d::D) where {P,T<:$SCT.HessianTracer,D<:$SCT.Dual{P,T}}
             x = $SCT.primal(d)
             p_out = $M.$op(x)
@@ -88,21 +94,6 @@ function overload_hessian_2_to_1(M, op)
                 $SCT.is_crossder_zero_global($M.$op),
             )
         end
-        function $M.$op(dx::D, dy::D) where {P,T<:$SCT.HessianTracer,D<:$SCT.Dual{P,T}}
-            x = $SCT.primal(dx)
-            y = $SCT.primal(dy)
-            p_out = $M.$op(x, y)
-            t_out = $SCT.hessian_tracer_2_to_1(
-                $SCT.tracer(dx),
-                $SCT.tracer(dy),
-                $SCT.is_firstder_arg1_zero_local($M.$op, x, y),
-                $SCT.is_seconder_arg1_zero_local($M.$op, x, y),
-                $SCT.is_firstder_arg2_zero_local($M.$op, x, y),
-                $SCT.is_seconder_arg2_zero_local($M.$op, x, y),
-                $SCT.is_crossder_zero_local($M.$op, x, y),
-            )
-            return $SCT.Dual(p_out, t_out)
-        end
 
         function $M.$op(tx::$SCT.HessianTracer, y::Real)
             return $SCT.hessian_tracer_1_to_1(
@@ -111,13 +102,35 @@ function overload_hessian_2_to_1(M, op)
                 $SCT.is_seconder_arg1_zero_global($M.$op),
             )
         end
+
         function $M.$op(x::Real, ty::$SCT.HessianTracer)
             return $SCT.hessian_tracer_1_to_1(
                 ty,
                 $SCT.is_firstder_arg2_zero_global($M.$op),
                 $SCT.is_seconder_arg2_zero_global($M.$op),
             )
         end
+    end
+end
+
+function overload_hessian_2_to_1_dual(M, op)
+    SCT = SparseConnectivityTracer
+    return quote
+        function $M.$op(dx::D, dy::D) where {P,T<:$SCT.HessianTracer,D<:$SCT.Dual{P,T}}
+            x = $SCT.primal(dx)
+            y = $SCT.primal(dy)
+            p_out = $M.$op(x, y)
+            t_out = $SCT.hessian_tracer_2_to_1(
+                $SCT.tracer(dx),
+                $SCT.tracer(dy),
+                $SCT.is_firstder_arg1_zero_local($M.$op, x, y),
+                $SCT.is_seconder_arg1_zero_local($M.$op, x, y),
+                $SCT.is_firstder_arg2_zero_local($M.$op, x, y),
+                $SCT.is_seconder_arg2_zero_local($M.$op, x, y),
+                $SCT.is_crossder_zero_local($M.$op, x, y),
+            )
+            return $SCT.Dual(p_out, t_out)
+        end
 
         function $M.$op(dx::D, y::Real) where {P,T<:$SCT.HessianTracer,D<:$SCT.Dual{P,T}}
             x = $SCT.primal(dx)
@@ -129,6 +142,7 @@ function overload_hessian_2_to_1(M, op)
             )
             return $SCT.Dual(p_out, t_out)
         end
+
         function $M.$op(x::Real, dy::D) where {P,T<:$SCT.HessianTracer,D<:$SCT.Dual{P,T}}
             y = $SCT.primal(dy)
             p_out = $M.$op(x, y)
@@ -168,7 +182,12 @@ function overload_hessian_1_to_2(M, op)
                 $SCT.is_seconder_out2_zero_global($M.$op),
             )
         end
+    end
+end
 
+function overload_hessian_1_to_2_dual(M, op)
+    SCT = SparseConnectivityTracer
+    return quote
         function $M.$op(d::D) where {P,T<:$SCT.HessianTracer,D<:$SCT.Dual{P,T}}
             x = $SCT.primal(d)
             p1_out, p2_out = $M.$op(x)