Loosen type rules for Mooncake

Author: kshyatt

ext/VectorInterfaceMooncakeExt.jl

@@ -23,8 +23,8 @@ _needs_tangent(::Type{T}) where {T <: Number} =
 
 # scale
 # -----
-@is_primitive DefaultCtx Tuple{typeof(scale!), AbstractArray, Number}
-function Mooncake.rrule!!(::CoDual{typeof(scale!)}, C_ΔC::CoDual{<:AbstractArray}, α_Δα::CoDual{<:Number})
+@is_primitive DefaultCtx Tuple{typeof(scale!), Any, Number}
+function Mooncake.rrule!!(::CoDual{typeof(scale!)}, C_ΔC::CoDual, α_Δα::CoDual{<:Number})
     # prepare arguments
     C, ΔC = arrayify(C_ΔC)
     α = primal(α_Δα)
@@ -43,7 +43,7 @@ function Mooncake.rrule!!(::CoDual{typeof(scale!)}, C_ΔC::CoDual{<:AbstractArra
     return C_ΔC, scale_pullback
 end
 
-function Mooncake.frule!!(::Dual{typeof(scale!)}, C_ΔC::Dual{<:AbstractArray}, α_Δα::Dual{<:Number})
+function Mooncake.frule!!(::Dual{typeof(scale!)}, C_ΔC::Dual, α_Δα::Dual{<:Number})
     # prepare arguments
     C, ΔC = arrayify(C_ΔC)
     α, Δα = extract(α_Δα)
@@ -58,9 +58,9 @@ function Mooncake.frule!!(::Dual{typeof(scale!)}, C_ΔC::Dual{<:AbstractArray},
     return C_ΔC
 end
 
-@is_primitive DefaultCtx Tuple{typeof(scale!), AbstractArray, AbstractArray, Number}
+@is_primitive DefaultCtx Tuple{typeof(scale!), Any, Any, Number}
 
-function Mooncake.rrule!!(::CoDual{typeof(scale!)}, C_ΔC::CoDual{<:AbstractArray}, A_ΔA::CoDual{<:AbstractArray}, α_Δα::CoDual{<:Number})
+function Mooncake.rrule!!(::CoDual{typeof(scale!)}, C_ΔC::CoDual, A_ΔA::CoDual, α_Δα::CoDual{<:Number})
     # prepare arguments
     C, ΔC = arrayify(C_ΔC)
     A, ΔA = arrayify(A_ΔA)
@@ -81,7 +81,7 @@ function Mooncake.rrule!!(::CoDual{typeof(scale!)}, C_ΔC::CoDual{<:AbstractArra
     return C_ΔC, scale_pullback
 end
 
-function Mooncake.frule!!(::Dual{typeof(scale!)}, C_ΔC::Dual{<:AbstractArray}, A_ΔA::Dual{<:AbstractArray}, α_Δα::Dual{<:Number})
+function Mooncake.frule!!(::Dual{typeof(scale!)}, C_ΔC::Dual, A_ΔA::Dual, α_Δα::Dual{<:Number})
     # prepare arguments
     C, ΔC = arrayify(C_ΔC)
     A, ΔA = arrayify(A_ΔA)
@@ -96,9 +96,9 @@ end
 # add
 # ---
 
-@is_primitive DefaultCtx Tuple{typeof(add!), AbstractArray, AbstractArray, Number, Number}
+@is_primitive DefaultCtx Tuple{typeof(add!), Any, Any, Number, Number}
 
-function Mooncake.rrule!!(::CoDual{typeof(add!)}, C_ΔC::CoDual{<:AbstractArray}, A_ΔA::CoDual{<:AbstractArray}, α_Δα::CoDual{<:Number}, β_Δβ::CoDual{<:Number})
+function Mooncake.rrule!!(::CoDual{typeof(add!)}, C_ΔC::CoDual, A_ΔA::CoDual, α_Δα::CoDual{<:Number}, β_Δβ::CoDual{<:Number})
     # prepare arguments
     C, ΔC = arrayify(C_ΔC)
     A, ΔA = arrayify(A_ΔA)
@@ -123,7 +123,7 @@ function Mooncake.rrule!!(::CoDual{typeof(add!)}, C_ΔC::CoDual{<:AbstractArray}
     return C_ΔC, add_pullback
 end
 
-function Mooncake.frule!!(::Dual{typeof(add!)}, C_ΔC::Dual{<:AbstractArray}, A_ΔA::Dual{<:AbstractArray}, α_Δα::Dual{<:Number}, β_Δβ::Dual{<:Number})
+function Mooncake.frule!!(::Dual{typeof(add!)}, C_ΔC::Dual, A_ΔA::Dual, α_Δα::Dual{<:Number}, β_Δβ::Dual{<:Number})
     # prepare arguments
     C, ΔC = arrayify(C_ΔC)
     A, ΔA = arrayify(A_ΔA)
@@ -140,9 +140,9 @@ end
 # inner
 # -----
 
-@is_primitive DefaultCtx Tuple{typeof(inner), AbstractArray, AbstractArray}
+@is_primitive DefaultCtx Tuple{typeof(inner), Any, Any}
 
-function Mooncake.rrule!!(::CoDual{typeof(inner)}, A_ΔA::CoDual{<:AbstractArray}, B_ΔB::CoDual{<:AbstractArray})
+function Mooncake.rrule!!(::CoDual{typeof(inner)}, A_ΔA::CoDual, B_ΔB::CoDual)
     # prepare arguments
     A, ΔA = arrayify(A_ΔA)
     B, ΔB = arrayify(B_ΔB)
@@ -159,7 +159,7 @@ function Mooncake.rrule!!(::CoDual{typeof(inner)}, A_ΔA::CoDual{<:AbstractArray
     return CoDual(s, NoFData()), inner_pullback
 end
 
-function Mooncake.frule!!(::Dual{typeof(inner)}, A_ΔA::Dual{<:AbstractArray}, B_ΔB::Dual{<:AbstractArray})
+function Mooncake.frule!!(::Dual{typeof(inner)}, A_ΔA::Dual, B_ΔB::Dual)
     # prepare arguments
     A, ΔA = arrayify(A_ΔA)
     B, ΔB = arrayify(B_ΔB)