 
   static __device__ void compute(
       IndexType linearId, const int dim,
-      const SimpleTensorInfo<IndexType>& index, IndexType* indexOffset,
-      const SimpleTensorInfo<IndexType>& t2, IndexType* t2Offset) {
+      const TensorInfo<IndexType>& index, IndexType* indexOffset,
+      const TensorInfo<IndexType>& t2, IndexType* t2Offset) {
     for (int d = index.dims - 1; d >= 0; d--) {
       IndexType curDimIndex = linearId % index.sizes[d];
       *indexOffset += curDimIndex * index.strides[d];

        break;
    }
  } else {
    SimpleTensorInfo<unsigned long> tensorInfo(state, tensor);
    SimpleTensorInfo<unsigned long> srcInfo(state, src);
    SimpleTensorInfo<unsigned long> indexInfo(state, index);

    RUN(unsigned long, -1)
  }

struct IndexToScatterGatherOffsets {
  static __device__ void compute(
      IndexType linearId, const int dim,
      const TensorInfo<IndexType>& index, IndexType* indexOffset,
      const TensorInfo<IndexType>& t1, IndexType* t1Offset,
      const TensorInfo<IndexType>& t2, IndexType* t2Offset) {
    for (int d = Dims - 1; d >= 0; d--) {
      IndexType curDimIndex = linearId % index.sizes[d];
      *indexOffset += curDimIndex * index.strides[d];
