initialize endMark

source/sample/transformer/T2TPredictor.cpp

@@ -103,9 +103,11 @@ void T2TPredictor::Create(T2TModel * model, XTensor * top, const XTensor * input
 
     InitTensor(&state->probPath, input->order, dims, X_FLOAT);
     InitTensor(&state->nstep, input->order, dims, X_FLOAT);
+    InitTensor(&state->endMark, input->order, dims, X_FLOAT);
 
     state->probPath.SetZeroAll();
     state->nstep.SetZeroAll();
+    state->endMark.SetZeroAll();
 }
 
 /* 

source/sample/transformer/T2TSearch.cpp

@@ -158,7 +158,6 @@ void T2TSearch::Score(T2TStateBundle * prev, T2TStateBundle * beam)
     _SumDim(&prob, &probPathPrev, &score, 0);
 
     InitTensor(&len, &lenPrev);
-    InitTensor(&lp, &lenPrev);
 
     _ScaleAndShift(&lenPrev, &len, 1.0F, 1.0F);
 
@@ -170,18 +169,21 @@ void T2TSearch::Score(T2TStateBundle * prev, T2TStateBundle * beam)
     /* score = log-prob/lp */
     _DivDim(&score, &lp, &score, 0);
 
-    prob.Reshape(order, dims);
-    score.Reshape(order, dims);
-    probPathPrev.Reshape(order - 1, dims);
-    lp.Reshape(order - 1, dims);
-
     InitTensor(&mask, &prev->endMark);
     CopyValues(prev->endMark, mask);
     _ScaleAndShiftMe(&mask, -1e9F);
+    
+    mask.Reshape(mask.unitNum);
 
     /* mask the completed hypotheses so that they cannot 
        be involved in further sorting and beam search. */
-    _Sum(&score, &mask, &score);
+    _SumDim(&score, &mask, &score, 0);
+    
+    prob.Reshape(order, dims);
+    score.Reshape(order, dims);
+    probPathPrev.Reshape(order - 1, dims);
+    lp.Reshape(order - 1, dims);
+    mask.Reshape(order -1 , dims);
 }
 
 /* 
@@ -225,7 +227,7 @@ void T2TSearch::Generate(T2TStateBundle * beam)
     score.Reshape(order, dimsBeam);
     
     /* keep the most promissing candidates in the beam */
-    TopK(score, scoreTopK, index, 0, beamSize);
+    TopK(score, scoreTopK, index, -1, beamSize);
     
     CopyValues(scoreTopK, preID);
     

source/tensor/core/sort/Sort.cpp

@@ -19,6 +19,7 @@
 * $Created by: XIAO Tong (email: xiaotong@mail.neu.edu.cn) 2018-04-24
 */
 
+#include <math.h>
 #include "../../XTensor.h"
 #include "../movement/CopyValues.h"
 #include "../../XUtility.h"
@@ -37,6 +38,8 @@ sort the tensor along a given dimension
 */
 void _Sort(const XTensor * a, XTensor * b, XTensor * index, int dim)
 {
+    dim = MODX(dim, a->order);
+    
     CheckNTErrors((XTensor::IsSameShaped(a, b)), "Input tensors should have the same type!");
     CheckNTErrors((dim >= 0 && dim < a->order), "Incorrect dimension specified!");
     CheckNTErrors((a->order == index->order), "Unmatched input tensors!");
@@ -105,6 +108,8 @@ make a new tensor to keep the result and return it
 */
 void Sort(XTensor & a, XTensor & b, XTensor & index, int dim)
 {
+    dim = MODX(dim, a.order);
+    
     /* call _Negate function */
     _Sort(&a, &b, &index, dim);
     

source/tensor/core/sort/TopK.cpp

@@ -19,6 +19,7 @@
 * $Created by: XIAO Tong (email: xiaotong@mail.neu.edu.cn) 2018-04-24
 */
 
+#include <math.h>
 #include "../../XTensor.h"
 #include "../../XName.h"
 #include "TopK.h"
@@ -37,6 +38,8 @@ get the top-k items along a given dimension
 */
 void _TopK(const XTensor * a, XTensor * b, XTensor * index, int dim, int k)
 {
+    dim = MODX(dim, a->order);
+    
     CheckNTErrors(a->unitSize == b->unitSize, "Unmatched input tensors!");
     CheckNTErrors(a->order == b->order, "Unmatched input tensors!");
     CheckNTErrors(index == NULL || a->order == index->order, "Unmatched input tensors!");
@@ -117,6 +120,8 @@ get the top-k items along a given dimension
 */
 void TopK(XTensor &a, XTensor &b, XTensor &index, int dim, int k)
 {
+    dim = MODX(dim, a.order);
+    
     if(a.dimSize[dim] <= k)
         _Sort(&a, &b, &index, dim);
     else