Add adjoint for complex number

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, batched_mul, batched_transpose
+  softmax, logsoftmax, maxpool, meanpool, batched_mul, batched_transpose, batched_adjoint
 
 include("numeric.jl")
 include("activation.jl")

src/linalg.jl

@@ -26,17 +26,27 @@ for (gemm, elty) in ((:dgemm_,:Float64), (:sgemm_,:Float32))
                   transA, transB, M, N, K,
                   alpha, A, lda, B, ldb, beta, C, ldc)
         end
+    end
+end
 
-        ##  borrow BatchedRoutines.jl
-        # batched gemm for 3d-array
-        # C[:,:,i] := alpha*op(A[:,:,i])*op(B[:,:,i]) + beta*C[:,:,i], where:
-        # i is the specific batch number,
-        # op(X) is one of op(X) = X, or op(X) = XT, or op(X) = XH,
-        # alpha and beta are scalars,
-        # A, B and C are 3d Array:
-        # op(A) is an m-by-k-by-b 3d Array,
-        # op(B) is a k-by-n-by-b 3d Array,
-        # C is an m-by-n-by-b 3d Array.
+
+##  borrow BatchedRoutines.jl
+# batched gemm for 3d-array
+# C[:,:,i] := alpha*op(A[:,:,i])*op(B[:,:,i]) + beta*C[:,:,i], where:
+# i is the specific batch number,
+# op(X) is one of op(X) = X, or op(X) = XT, or op(X) = XH,
+# alpha and beta are scalars,
+# A, B and C are 3d Array:
+# op(A) is an m-by-k-by-b 3d Array,
+# op(B) is a k-by-n-by-b 3d Array,
+# C is an m-by-n-by-b 3d Array.
+
+for (gemm, elty) in
+        ((:dgemm_,:Float64),
+         (:sgemm_,:Float32),
+         (:zgemm_,:ComplexF64),
+         (:cgemm_,:ComplexF32))
+    @eval begin
         function batched_gemm!(transA::AbstractChar,
                                transB::AbstractChar,
                                alpha::($elty),

test/batchedmul.jl

@@ -10,6 +10,18 @@ function bmm_test(a,b; transA = false, transB = false)
     cat(c...; dims = 3)
 end
 
+function bmm_adjtest(a,b; adjA = false, adjB = false)
+    bs = size(a,3)
+    c = []
+    for i = 1:bs
+        ai = adjA ? adjoint(a[:,:,i]) : a[:,:,i]
+        bi = adjB ? adjoint(b[:,:,i]) : b[:,:,i]
+        push!(c, ai*bi)
+    end
+
+    cat(c...; dims = 3)
+end
+
 @testset "Batched Matrix Multiplication" begin
     A = randn(7,5,3)
     B = randn(5,7,3)
@@ -19,4 +31,14 @@ end
     @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)
+
+
+    cA = randn(Complex{Float64}, 7,5,3)
+    cB = randn(Complex{Float64}, 5,7,3)
+    cC = randn(Complex{Float64}, 7,6,3)
+
+    @test batched_mul(cA, cB) == bmm_adjtest(cA, cB)
+    @test batched_mul(batched_adjoint(cA), batched_adjoint(cB)) == bmm_adjtest(cA, cB; adjA = true, adjB = true)
+    @test batched_mul(batched_adjoint(cA), cC) == bmm_adjtest(cA, cC; adjA = true)
+    @test batched_mul(cA, batched_adjoint(cA)) == bmm_adjtest(cA, cA; adjB = true)
 end