Add some custom rrule

Author: avik-pal

src/layers.jl

@@ -335,7 +335,7 @@ If `init` is not passed, it creates a MultiScale Regularized Deep Equilibrium Ne
 """
 function MultiScaleSkipDeepEquilibriumNetwork(main_layers::Tuple, mapping_layers::Matrix,
         post_fuse_layer::Union{Nothing, Tuple}, init::Tuple, solver, scales; kwargs...)
-    init = Chain(Parallel(nothing, init...), x -> mapreduce(__flatten, vcat, x))
+    init = Chain(Parallel(nothing, init...), __flatten_vcat)
     return MultiScaleDeepEquilibriumNetwork(main_layers, mapping_layers, post_fuse_layer,
         solver, scales; init, kwargs...)
 end
@@ -393,6 +393,6 @@ end
         u, x = z
         u_ = __split_and_reshape(u, m.split_idxs, m.scales)
         u_res, st = Lux.apply(m.model, ($(inputs...),), ps, st)
-        return mapreduce(__flatten, vcat, u_res), st
+        return __flatten_vcat(u_res), st
     end
 end

src/utils.jl

@@ -11,10 +11,31 @@ end
 __split_and_reshape(x::AbstractMatrix, ::Nothing, ::Nothing) = x
 __split_and_reshape(x::AbstractArray, ::Nothing, ::Nothing) = x
 
+function __split_and_reshape(y::AbstractMatrix, x)
+    szs = [prod(size(xᵢ)[1:(end - 1)]) for xᵢ in x]
+    counters = vcat(0, cumsum(szs)[1:(end - 1)])
+    return map((sz, c, xᵢ) -> reshape(view(y, (c + 1):(c + sz), :), size(xᵢ)),
+        szs, counters, x)
+end
+
 @inline __flatten(x::AbstractVector) = reshape(x, length(x), 1)
 @inline __flatten(x::AbstractMatrix) = x
 @inline __flatten(x::AbstractArray) = reshape(x, :, size(x, ndims(x)))
 
+@inline __flatten_vcat(x) = mapreduce(__flatten, vcat, x)
+
+function CRC.rrule(::typeof(__flatten_vcat), x)
+    y = __flatten_vcat(x)
+    projects = CRC.ProjectTo.(x)
+    function ∇__flatten_vcat(∂y)
+        ∂y isa CRC.NoTangent && return (CRC.NoTangent(), CRC.NoTangent())
+        ∂x = __split_and_reshape(∂y, x)
+        ∂x = map((∂xᵢ, project) -> project(∂xᵢ), ∂x, projects)
+        return CRC.NoTangent(), ∂x
+    end
+    return y, ∇__flatten_vcat
+end
+
 @inline __check_unrolled_mode(::Val{d}) where {d} = Val(d ≥ 1)
 @inline __check_unrolled_mode(st::NamedTuple) = __check_unrolled_mode(st.fixed_depth)