bug fixes in matrix multiplication

source/network/XBackwardMath.cpp

@@ -259,10 +259,52 @@ void XMathGrad::GradMatrixMul(XTensor * node)
     XNoder::MakeGrad(a);
     XNoder::MakeGrad(b);
 
+    XTensor * c = node;
     XTensor * dedc = node->grad;
     XTensor * deda = a->grad;
     XTensor * dedb = b->grad;
     
+    if(deda->order == 2 && dedb->order == 2)
+        GradMatrixMul(a, deda, transA, b, dedb, transB, dedc, alpha);
+    else if(transA == X_NOTRANS && deda->order > 2 && dedb->order == 2){
+        int orderBackupA = a->order;
+        int orderBackupC = c->order;
+        int dimsBackupA[MAX_TENSOR_DIM_NUM];
+        int dimsBackupC[MAX_TENSOR_DIM_NUM];
+        memcpy(dimsBackupA, a->dimSize, sizeof(int) * a->order);
+        memcpy(dimsBackupC, c->dimSize, sizeof(int) * c->order);
+        int dimsA[2] = {a->unitNum/a->GetDim(-1), a->GetDim(-1)};
+        int dimsC[2] = {c->unitNum/c->GetDim(-1), c->GetDim(-1)};
+        a->Reshape(2, dimsA);
+        c->Reshape(2, dimsC);
+        deda->Reshape(2, dimsA);
+        dedc->Reshape(2, dimsC);
+        GradMatrixMul(a, deda, transA, b, dedb, transB, dedc, alpha);
+        a->Reshape(orderBackupA, dimsBackupA);
+        c->Reshape(orderBackupC, dimsBackupC);
+        deda->Reshape(orderBackupA, dimsBackupA);
+        dedc->Reshape(orderBackupC, dimsBackupC);
+    }
+    else{
+        ShowNTErrors("TODO!");
+    }
+
+    node->visitMark = NODE_FINISHED;
+}
+    
+/*
+gradient for matrix multiply: c = matmul(a, b) * \alpha
+>> a - as it is
+>> deda - dE/da
+>> b - as it is
+>> dedb - dE/db
+>> dedc - dE/dc
+>> alpha - the scalar
+*/
+void XMathGrad::GradMatrixMul(XTensor * a, XTensor * deda, MATRIX_TRANS_TYPE transA,
+                              XTensor * b, XTensor * dedb, MATRIX_TRANS_TYPE transB,
+                              XTensor * dedc, DTYPE alpha)
+{
     /* c = a * b * \alpha */
     if(transA == X_NOTRANS && transB == X_NOTRANS){
         
@@ -302,8 +344,6 @@ void XMathGrad::GradMatrixMul(XTensor * node)
         /* dE/db = a * dE/dc * \alpha */
         _MatrixMul(a, X_NOTRANS, dedc, X_NOTRANS, dedb, alpha, 1.0F);
     }
-
-    node->visitMark = NODE_FINISHED;
 }
 
 /*

source/network/XBackwardMath.h

@@ -57,6 +57,12 @@ private:
     static
     void GradMatrixMul(XTensor * node);
     
+    /* gradient for matrix multiply: c = matmul(a, b) */
+    static
+    void GradMatrixMul(XTensor * a, XTensor * deda, MATRIX_TRANS_TYPE transA,
+                       XTensor * b, XTensor * dedb, MATRIX_TRANS_TYPE transB,
+                       XTensor * dedc, DTYPE alpha);
+
     /* gradient for log: c =  log(a) */
     static
     void GradLog(XTensor * node);

source/sample/transformer/T2TUtility.cpp

@@ -26,7 +26,7 @@
 namespace transformer
 {
 
-void LoadParamString(int argc, const char ** argv, const char * name, char * p, char * defaultP)
+void LoadParamString(int argc, const char ** argv, const char * name, char * p, const char * defaultP)
 {
     char vname[128];
     vname[0] = '-';

source/sample/transformer/T2TUtility.h

@@ -28,7 +28,7 @@ namespace transformer
 {
 
 /* load arguments */
-void LoadParamString(int argc, const char ** argv, const char * name, char * p, char * defaultP);
+void LoadParamString(int argc, const char ** argv, const char * name, char * p, const char * defaultP);
 void LoadParamInt(int argc, const char ** argv, const char * name, int * p, int defaultP);
 void LoadParamBool(int argc, const char ** argv, const char * name, bool * p, bool defaultP);
 void LoadParamFloat(int argc, const char ** argv, const char * name, float * p, float defaultP);

source/tensor/core/arithmetic/MatrixMul.cpp

@@ -53,11 +53,29 @@ void _MatrixMul(const XTensor * a, MATRIX_TRANS_TYPE transposedA,
                 const XTensor * b, MATRIX_TRANS_TYPE transposedB,
                 XTensor * c, DTYPE alpha, DTYPE beta, XPRunner * parallelRunner)
 {
-    CheckNTErrors((a && b && c), "Empty input tensors!");
-    CheckNTErrors((a->dataType == b->dataType && a->dataType == c->dataType),
+    CheckNTErrors(a && b && c, "Empty input tensors!");
+    CheckNTErrors(a->dataType == b->dataType && a->dataType == c->dataType,
                   "Input tensors should have the same data type!");
-    CheckNTErrors((a->order >= 2 && b->order >= 2 && c->order >= 2),
+    CheckNTErrors(a->order >= 2 && b->order >= 2 && c->order >= 2,
                   "Input tensors must have a order >= 2!");
+    CheckNTErrors(c->order == a->order + b->order - 2, "wrong tensor order")
+    
+    /* we transform a higher order tensor to a matrix to kill the number
+       of calls of matrix multiplication */
+    if(transposedA == X_NOTRANS && a->order > 2 && b->order == 2){
+        int ncolA = a->dimSize[a->order - 1];
+        int ncolC = c->dimSize[c->order - 1];
+        XTensor * a2 = NewTensor2D(a->unitNum/ncolA, -ncolA, a->dataType, a->devID, a->mem);
+        XTensor * c2 = NewTensor2D(c->unitNum/ncolC, -ncolC, c->dataType, c->devID, c->mem);
+        a2->data = a->data;
+        c2->data = c->data;
+        _MatrixMul2D(a2, transposedA, b, transposedB, c2, alpha, beta, parallelRunner);
+        a2->data = NULL;
+        c2->data = NULL;
+        delete a2;
+        delete c2;
+        return;
+    }
 
     int an = transposedA == X_TRANS ? a->dimSizeRDI[0] : a->dimSizeRDI[1];
     int am = transposedA == X_TRANS ? a->dimSizeRDI[1] : a->dimSizeRDI[0];