gradient for adjoint bmm

Author: chengchingwen

src/batchedmul.jl

@@ -51,3 +51,16 @@ end
 function ∇batched_mul(Δ::AbstractArray{T, 3}, A::BatchedTranspose{T, <: AbstractArray{T, 3}}, B::BatchedTranspose{T, <: AbstractArray{T, 3}}) where T
     (batched_mul(batched_transpose(Δ), batched_transpose(B)), batched_mul(batched_transpose(A), batched_transpose(Δ)))
 end
+
+function ∇batched_mul(Δ::AbstractArray{T, 3}, A::BatchedAdjoint{T, <: AbstractArray{T, 3}}, B::AbstractArray{T, 3}) where T
+    (batched_mul(Δ, batched_adjoint(B)), batched_mul(A, Δ))
+end
+
+function ∇batched_mul(Δ::AbstractArray{T, 3}, A::AbstractArray{T, 3}, B::BatchedAdjoint{T, <: AbstractArray{T, 3}}) where T
+    (batched_mul(Δ, B), batched_mul(batched_adjoint(A), Δ))
+end
+
+function ∇batched_mul(Δ::AbstractArray{T, 3}, A::BatchedAdjoint{T, <: AbstractArray{T, 3}}, B::BatchedAdjoint{T, <: AbstractArray{T, 3}}) where T
+    (batched_mul(batched_adjoint(Δ), batched_adjoint(B)), batched_mul(batched_adjoint(A), batched_adjoint(Δ)))
+end
+