port batched transpose type, remove transpose keyword

Author: chengchingwen

src/NNlib.jl

@@ -3,7 +3,7 @@ module NNlib
 using Requires, Libdl
 
 export σ, sigmoid, relu, leakyrelu, elu, gelu, swish, selu, softplus, softsign, logσ, logsigmoid,
-  softmax, logsoftmax, maxpool, meanpool, batchedmul
+  softmax, logsoftmax, maxpool, meanpool, batched_mul, batched_transpose
 
 include("numeric.jl")
 include("activation.jl")

src/batchedadjtrans.jl

@@ -0,0 +1,67 @@
+using LinearAlgebra
+import Base: -
+
+"""
+    BatchedTranspose{T, N, S} <: AbstractBatchedMatrix{T, N}
+Batched transpose. Transpose a batch of matrix.
+"""
+struct BatchedTranspose{T, S} <: AbstractArray{T, 3}
+    parent::S
+    BatchedTranspose{T, S}(X::S) where {T, S} = new{T, S}(X)
+end
+
+"""
+    batched_transpose(A)
+Lazy batched transpose.
+"""
+batched_transpose(A::AbstractArray{T}) where T = BatchedTranspose(A)
+
+
+"""
+    BatchedAdjoint{T, N, S} <: AbstractBatchedMatrix{T, N}
+Batched ajoint. Transpose a batch of matrix.
+"""
+struct BatchedAdjoint{T, S} <: AbstractArray{T, 3}
+    parent::S
+    BatchedAdjoint{T, S}(X::S) where {T, S} = new{T, S}(X)
+end
+
+"""
+    batched_adjoint(A)
+Lazy batched adjoint.
+"""
+batched_adjoint(A::AbstractArray{T, 3}) where T = BatchedAdjoint(A)
+
+
+BatchedAdjoint(A) = BatchedAdjoint{Base.promote_op(adjoint,eltype(A)),typeof(A)}(A)
+BatchedTranspose(A) = BatchedTranspose{Base.promote_op(transpose,eltype(A)),typeof(A)}(A)
+
+
+const BatchedAdjOrTrans{T, S} = Union{BatchedTranspose{T, S}, BatchedAdjoint{T, S}}
+
+LinearAlgebra.wrapperop(A::BatchedAdjoint) = batched_adjoint
+LinearAlgebra.wrapperop(B::BatchedTranspose) = batched_transpose
+
+# AbstractArray Interface
+Base.length(A::BatchedAdjOrTrans) = length(A.parent)
+Base.size(m::BatchedAdjOrTrans) = (size(m.parent, 2), size(m.parent, 1), size(m.parent, 3))
+Base.axes(m::BatchedAdjOrTrans) = (axes(m.parent, 2), axes(m.parent, 1), axes(m.parent, 3))
+
+Base.IndexStyle(::Type{<:BatchedAdjOrTrans}) = IndexCartesian()
+Base.@propagate_inbounds Base.getindex(m::BatchedTranspose, i::Int, j::Int, k::Int) = getindex(m.parent, j, i, k)
+Base.@propagate_inbounds Base.getindex(m::BatchedAdjoint, i::Int, j::Int, k::Int) = adjoint(getindex(m.parent, j, i, k))
+Base.@propagate_inbounds Base.setindex!(m::BatchedAdjOrTrans, v, i::Int, j::Int, k::Int) = setindex!(m.parent, v, j, i, k)
+
+Base.similar(A::BatchedAdjOrTrans, T::Type, dims::Dims) = similar(A.parent, T, dims)
+Base.similar(A::BatchedAdjOrTrans, dims::Dims) = similar(A.parent, dims)
+Base.similar(A::BatchedAdjOrTrans, T::Type) = similar(A.parent, T, size(A))
+Base.similar(A::BatchedAdjOrTrans) = similar(A.parent, size(A))
+
+Base.parent(A::BatchedAdjOrTrans) = A.parent
+
+(-)(A::BatchedAdjoint)   = BatchedAdjoint(  -A.parent)
+(-)(A::BatchedTranspose) = BatchedTranspose(-A.parent)
+
+Base.copy(A::BatchedTranspose) = BatchedTranspose(copy(A.parent))
+Base.copy(A::BatchedAdjoint) = BatchedAdjoint(copy(A.parent))
+

src/batchedmul.jl

@@ -1,36 +1,53 @@
 # batch-wise matrix multiplication
 # wrapper for batched_gemm!
 
-function batchedmul(a::AbstractArray{T, 3}, b::AbstractArray{T, 3};
-                    transA::Bool = false, transB::Bool = false) where T
-    (bs = size(a, 3)) == size(b, 3) || error("batch size mismatch")
-    res = similar(a, size(a, transA ? 2 : 1), size(b, transB ? 1 : 2), bs)
-    batched_mul!(res, a, b; transA=transA, transB=transB)
-    return res
-end
+include("./batchedadjtrans.jl")
 
-function batched_mul!(C::AbstractArray{T, 3}, A::AbstractArray{T, 3}, B::AbstractArray{T, 3};
-                      transA::Bool = false, transB::Bool = false) where T
-    At = transA ? 'T' : 'N'
-    Bt = transB ? 'T' : 'N'
-    batched_gemm!(At, Bt, one(T), A, B, zero(T), C)
-    C
+function batched_mul(A::AbstractArray{T, 3}, B::AbstractArray{T, 3}) where T
+    size(A, 3) == size(B, 3) || throw(DimensionMismatch("batch size mismatch"))
+    batched_mul!(similar(A, (size(A, 1), size(B, 2), size(A, 3))), A, B)
 end
 
-#gradient function for batchedmul
-function ∇batchedmul(Δ::AbstractArray{T, 3}, a::AbstractArray{T, 3}, b::AbstractArray{T, 3};
-                     transA::Bool = false, transB::Bool = false) where T
-    if transA
-        if transB
-            (batchedmul(b, Δ; transA=true, transB=true), batchedmul(Δ, a; transA=true, transB=true))
-        else
-            (batchedmul(b, Δ; transB=true), batchedmul(a, Δ))
-        end
-    else
-        if transB
-            (batchedmul(Δ, b), batchedmul(Δ, a; transA=true))
-        else
-            (batchedmul(Δ, b; transB=true), batchedmul(a, Δ; transA=true))
+"""
+    batched_mul!(C, A, B) -> C
+batched `mul!`.
+"""
+function batched_mul! end
+
+_unbatch(A) = A
+_unbatch(A::BatchedAdjOrTrans) = A.parent
+
+# bmm
+const _BATCHED_MATRIX_LIST = [
+        (:(AbstractArray{T, 3}), 'N'),
+        (:(BatchedTranspose{T, <:AbstractArray{T, 3}}), 'T'),
+        (:(BatchedAdjoint{T, <:AbstractArray{T, 3}}), 'C')
+]
+
+for (TA, transA) in _BATCHED_MATRIX_LIST, (TB, transB) in _BATCHED_MATRIX_LIST
+    @eval begin
+        function batched_mul!(C::AbstractArray{T, 3}, A::$TA, B::$TB) where T
+            batched_gemm!($transA, $transB, one(T), _unbatch(A), _unbatch(B), zero(T), C)
+            C
         end
+
+
     end
 end
+
+function ∇batched_mul(Δ::AbstractArray{T, 3}, A::AbstractArray{T, 3}, B::AbstractArray{T, 3}) where T
+    (batched_mul(Δ, batched_transpose(B)), batched_mul(batched_transpose(A), Δ))
+end
+
+
+function ∇batched_mul(Δ::AbstractArray{T, 3}, A::BatchedTranspose{T, <: AbstractArray{T, 3}}, B::AbstractArray{T, 3}) where T
+    (batched_mul(Δ, batched_transpose(B)), batched_mul(A, Δ))
+end
+
+function ∇batched_mul(Δ::AbstractArray{T, 3}, A::AbstractArray{T, 3}, B::BatchedTranspose{T, <: AbstractArray{T, 3}}) where T
+    (batched_mul(Δ, B), batched_mul(batched_transpose(A), Δ))
+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

test/batchedmul.jl

@@ -15,8 +15,8 @@ end
     B = randn(5,7,3)
     C = randn(7,6,3)
 
-    @test batchedmul(A, B) == bmm_test(A, B)
-    @test batchedmul(A, B; transA = true, transB = true) == bmm_test(A, B; transA = true, transB = true)
-    @test batchedmul(A, C; transA = true) == bmm_test(A, C; transA = true)
-    @test batchedmul(A, A; transB = true) == bmm_test(A, A; transB = true)
+    @test batched_mul(A, B) == bmm_test(A, B)
+    @test batched_mul(batched_transpose(A), batched_transpose(B)) == bmm_test(A, B; transA = true, transB = true)
+    @test batched_mul(batched_transpose(A), C) == bmm_test(A, C; transA = true)
+    @test batched_mul(A, batched_transpose(A)) == bmm_test(A, A; transB = true)
 end