Override chainrules rrule

LuxDL · Mar 28, 2024 · d0ce078 · d0ce078
1 parent 1911ccf
commit d0ce078
Showing 1 changed file with 8 additions and 0 deletions.
diff --git a/src/chainrules.jl b/src/chainrules.jl
@@ -115,6 +115,14 @@ function CRC.rrule(::typeof(*), X::UniformBlockDiagonalMatrix{<:Union{Real, Comp
     return X * Y, ∇times
 end
 
+function CRC.rrule(cfg::CRC.RuleConfig{>:CRC.HasReverseMode}, ::typeof(*),
+        X::AbstractMatrix{<:Union{Real, Complex}},
+        Y::UniformBlockDiagonalMatrix{<:Union{Real, Complex}})
+    _f = @closure (x, y) -> dropdims(
+        batched_mul(reshape(x, :, 1, nbatches(x)), y.data); dims=1)
+    return CRC.rrule_via_ad(cfg, _f, X, Y)
+end
+
 # constructor
 function CRC.rrule(::Type{<:UniformBlockDiagonalMatrix}, data)
     function ∇UniformBlockDiagonalMatrix(Δ)