select a range of data in a tensor

source/core/Select.cpp

+/* NiuTrans.Tensor - an open-source tensor library
+* Copyright (C) 2018, Natural Language Processing Lab, Northestern University.
+* All rights reserved.
+*
+* Licensed under the Apache License, Version 2.0 (the "License");
+* you may not use this file except in compliance with the License.
+* You may obtain a copy of the License at
+*
+*   http://www.apache.org/licenses/LICENSE-2.0
+*
+* Unless required by applicable law or agreed to in writing, software
+* distributed under the License is distributed on an "AS IS" BASIS,
+* WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+* See the License for the specific language governing permissions and
+* limitations under the License.
+*/
+
+/*
+* $Created by: XIAO Tong (email: xiaotong@mail.neu.edu.cn) 2018-07-04
+*/
+
+#include "Select.h"
+#include "../XUtility.h"
+
+namespace nts{ // namespace nts(NiuTrans.Tensor)
+
+/* 
+generate a tensor with seleccted data in range[low,high] along the given dimension 
+c = select(a) 
+>> a - input tensor
+>> dim - the dimension along with which we do the job
+>> low - lower bound
+>> high - higher bound.
+          Note that range [1,3] means that we select 1 and 2.
+>> c - result tensor
+*/
+void SelectRange(XTensor * a, int dim, int low, int high, XTensor * c)
+{
+    CheckNTErrors(a != NULL && c != NULL, "empty tensors!");
+    CheckNTErrors(a->order == c->order, "The input and output tensors must in the same order!");
+    CheckNTErrors(dim >= 0 && dim < a->order, "The input dimension is out of bounds!");
+    CheckNTErrors(a->dataType == c->dataType, "The tensor must be of the same data type!");
+    
+    if(low >= high)
+        return;
+
+    for(int i = 0; i < a->order; i++){
+        if(i == dim){
+            CheckNTErrors(low > 0 && low < a->dimSize[dim], "Illegal range specified!");
+            CheckNTErrors(high > 0 && high < a->dimSize[dim], "Illegal range specified!");
+        }
+        else{
+            CheckNTErrors(a->dimSize[i] == c->dimSize[i], "The size of the dimensions should be same!");
+        }
+    }
+
+    int stride = 1;
+    for(int i = 0; i < a->order; i++)
+        stride *= a->dimSizeRDI[i];
+
+    int blockSize = stride * (high - low) * a->unitSize;
+    int stepSizeS = stride * a->dimSize[dim] * a->unitSize;
+    int stepSizeT = stride * c->dimSize[dim] * a->unitSize;
+    char * s = (char*)a->data + stride * low * a->unitSize;
+    char * t = (char*)c->data;
+    for(int i = 0; i < high - low; i++){
+        XMemCopy(t, c->devID, s, a->devID, blockSize);
+        s += stepSizeS;
+        t += stepSizeT;
+    }
+}
+
+
+} // namespace nts(NiuTrans.Tensor)