Extend inplace mul

LuxDL · Mar 21, 2024 · 4facdbb · 4facdbb
1 parent 8037aa7
commit 4facdbb
Show file tree

Hide file tree

Showing 2 changed files with 9 additions and 1 deletion.
diff --git a/src/api.jl b/src/api.jl
@@ -55,7 +55,7 @@ TODO: Needs Documentation (take from NNlib.jl)
 """
 batched_mul(A, B) = _batched_mul(A, B)
 
-batched_mul!(C, A, B) = _batched_mul!(C, A, B)
+batched_mul!(C, A, B, α = true, β = false) = _batched_mul!(C, A, B, α, β)
 
 """
     batched_transpose(X::AbstractArray{T, 3}) where {T}

diff --git a/src/matrix.jl b/src/matrix.jl
@@ -282,6 +282,14 @@ function Base.:*(X::AbstractArray{T, 2}, Y::UniformBlockDiagonalMatrix) where {T
     return dropdims(C.data; dims=1)
 end
 
+function LinearAlgebra.mul!(A::AbstractMatrix, B::AbstractMatrix,
+        C::UniformBlockDiagonalMatrix, α::Number=true, β::Number=false)
+    A_ = reshape(A, 1, :, nbatches(A))
+    B_ = reshape(B, 1, :, nbatches(B))
+    batched_mul!(A_, B_, C.data, α, β)
+    return A
+end
+
 # LinearAlgebra
 abstract type AbstractBatchedMatrixFactorization{T} <: LinearAlgebra.Factorization{T} end