numpy,taco: adjust tensor shifting to shift less for higher order tensors

Author: rohany

numpy/util.py

@@ -171,7 +171,11 @@ def shiftLastMode(self, tensor):
             # resultValues[i] = data[i]
             # TODO (rohany): Temporarily use a constant as the value.
             resultValues[i] = 2
-            resultCoords[-1][i] = (resultCoords[-1][i] + 1) % tensor.shape[-1]
+            # For order 2 tensors, always shift the last coordinate. Otherwise, shift only coordinates
+            # that have even last coordinates. This ensures that there is at least some overlap
+            # between the original tensor and its shifted counter part.
+            if tensor.shape[-1] <= 0 or resultCoords[-1][i] % 2 == 0:
+                resultCoords[-1][i] = (resultCoords[-1][i] + 1) % tensor.shape[-1]
         return sparse.COO(resultCoords, resultValues, tensor.shape)
 
 # ScipyTensorShifter shifts all elements in the last mode

taco/bench.h

@@ -50,7 +50,12 @@ taco::Tensor<T> shiftLastMode(std::string name, taco::Tensor<T2> original) {
       coords[i] = value.first[i];
     }
     int lastMode = original.getOrder() - 1;
-    coords[lastMode] = (coords[lastMode] + 1) % original.getDimension(lastMode);
+    // For order 2 tensors, always shift the last coordinate. Otherwise, shift only coordinates
+    // that have even last coordinates. This ensures that there is at least some overlap
+    // between the original tensor and its shifted counter part.
+    if (original.getOrder() <= 2 || (coords[lastMode] % 2 == 0)) {
+      coords[lastMode] = (coords[lastMode] + 1) % original.getDimension(lastMode);
+    }
     // TODO (rohany): Temporarily use a constant value here.
     result.insert(coords, T(2));
   }