Make post_deq type-stable

Author: avik-pal

src/layers/chain.jl

@@ -27,22 +27,22 @@ function DEQChain(layers...)
         push!(encounter_deq ? post_deq : pre_deq, l)
     end
     @assert encounter_deq "No DEQ Layer in the Chain!!! Maybe you wanted to use Chain"
-    pre_deq = length(pre_deq) == 0 ? nothing : Chain(pre_deq...)
-    post_deq = length(post_deq) == 0 ? nothing : Chain(post_deq...)
+    pre_deq = length(pre_deq) == 0 ? NoOpLayer() : Chain(pre_deq...)
+    post_deq = length(post_deq) == 0 ? NoOpLayer() : Chain(post_deq...)
     return DEQChain(pre_deq, deq, post_deq)
 end
 
-function (deq::DEQChain{P1,D,P2})(x, ps::Union{ComponentArray,NamedTuple}, st::NamedTuple) where {P1,D,P2}
-    x1, st1 = if P1 == Nothing
-        x, st.pre_deq
-    else
-        deq.pre_deq(x, ps.pre_deq, st.pre_deq)
-    end
-    (x2, deq_soln), st2 = deq.deq(x1, ps.deq, st.deq)
-    x3, st3 = if P2 == Nothing
-        x2, st.post_deq
-    else
-        deq.post_deq(x2, ps.post_deq, st.post_deq)
-    end
+function get_deq_return_type(
+    deq::DEQChain{P1,<:Union{MultiScaleDeepEquilibriumNetwork,MultiScaleSkipDeepEquilibriumNetwork}}, ::T
+) where {P1,T}
+    return NTuple{length(deq.deq.scales),T}
+end
+get_deq_return_type(::DEQChain, ::T) where {T} = T
+
+function (deq::DEQChain)(x, ps::Union{ComponentArray,NamedTuple}, st::NamedTuple)
+    T = get_deq_return_type(deq, x)
+    x1, st1 = deq.pre_deq(x, ps.pre_deq, st.pre_deq)
+    (x2::T, deq_soln), st2 = deq.deq(x1, ps.deq, st.deq)
+    x3, st3 = deq.post_deq(x2, ps.post_deq, st.post_deq)
     return (x3, deq_soln), (pre_deq=st1, deq=st2, post_deq=st3)
 end